<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN" "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">
<html xmlns="http://www.w3.org/1999/xhtml">
<head>
<meta http-equiv="Content-Type" content="text/xhtml;charset=UTF-8"/>
<meta http-equiv="X-UA-Compatible" content="IE=9"/>
<meta name="generator" content="Doxygen 1.9.1"/>
<meta name="viewport" content="width=device-width, initial-scale=1"/>
<title>Eigen-unsupported: TensorContractionThreadPool.h Source File</title>
<link href="tabs.css" rel="stylesheet" type="text/css"/>
<script type="text/javascript" src="jquery.js"></script>
<script type="text/javascript" src="dynsections.js"></script>
<link href="navtree.css" rel="stylesheet" type="text/css"/>
<script type="text/javascript" src="resize.js"></script>
<script type="text/javascript" src="navtreedata.js"></script>
<script type="text/javascript" src="navtree.js"></script>
<link href="search/search.css" rel="stylesheet" type="text/css"/>
<script type="text/javascript" src="search/searchdata.js"></script>
<script type="text/javascript" src="search/search.js"></script>
<script type="text/javascript">
/* @license magnet:?xt=urn:btih:cf05388f2679ee054f2beb29a391d25f4e673ac3&amp;dn=gpl-2.0.txt GPL-v2 */
  $(document).ready(function() { init_search(); });
/* @license-end */
</script>
<script type="text/x-mathjax-config">
  MathJax.Hub.Config({
    extensions: ["tex2jax.js", "TeX/AMSmath.js", "TeX/AMSsymbols.js"],
    jax: ["input/TeX","output/HTML-CSS"],
});
</script>
<script type="text/javascript" async="async" src="https://cdn.mathjax.org/mathjax/latest/MathJax.js"></script>
<link href="doxygen.css" rel="stylesheet" type="text/css" />
<link href="eigendoxy.css" rel="stylesheet" type="text/css">
<!--  -->
<script type="text/javascript" src="eigen_navtree_hacks.js"></script>
</head>
<body>
<div id="top"><!-- do not remove this div, it is closed by doxygen! -->
<div id="titlearea">
<table cellspacing="0" cellpadding="0">
 <tbody>
 <tr style="height: 56px;">
  <td id="projectlogo"><img alt="Logo" src="Eigen_Silly_Professor_64x64.png"/></td>
  <td id="projectalign" style="padding-left: 0.5em;">
   <div id="projectname"><a href="http://eigen.tuxfamily.org">Eigen-unsupported</a>
   &#160;<span id="projectnumber">3.4.90 (git rev 67eeba6e720c5745abc77ae6c92ce0a44aa7b7ae)</span>
   </div>
  </td>
   <td>        <div id="MSearchBox" class="MSearchBoxInactive">
        <span class="left">
          <img id="MSearchSelect" src="search/mag_sel.svg"
               onmouseover="return searchBox.OnSearchSelectShow()"
               onmouseout="return searchBox.OnSearchSelectHide()"
               alt=""/>
          <input type="text" id="MSearchField" value="Search" accesskey="S"
               onfocus="searchBox.OnSearchFieldFocus(true)" 
               onblur="searchBox.OnSearchFieldFocus(false)" 
               onkeyup="searchBox.OnSearchFieldChange(event)"/>
          </span><span class="right">
            <a id="MSearchClose" href="javascript:searchBox.CloseResultsWindow()"><img id="MSearchCloseImg" border="0" src="search/close.svg" alt=""/></a>
          </span>
        </div>
</td>
 </tr>
 </tbody>
</table>
</div>
<!-- end header part -->
<!-- Generated by Doxygen 1.9.1 -->
<script type="text/javascript">
/* @license magnet:?xt=urn:btih:cf05388f2679ee054f2beb29a391d25f4e673ac3&amp;dn=gpl-2.0.txt GPL-v2 */
var searchBox = new SearchBox("searchBox", "search",false,'Search','.html');
/* @license-end */
</script>
</div><!-- top -->
<div id="side-nav" class="ui-resizable side-nav-resizable">
  <div id="nav-tree">
    <div id="nav-tree-contents">
      <div id="nav-sync" class="sync"></div>
    </div>
  </div>
  <div id="splitbar" style="-moz-user-select:none;" 
       class="ui-resizable-handle">
  </div>
</div>
<script type="text/javascript">
/* @license magnet:?xt=urn:btih:cf05388f2679ee054f2beb29a391d25f4e673ac3&amp;dn=gpl-2.0.txt GPL-v2 */
$(document).ready(function(){initNavTree('TensorContractionThreadPool_8h_source.html',''); initResizable(); });
/* @license-end */
</script>
<div id="doc-content">
<!-- window showing the filter options -->
<div id="MSearchSelectWindow"
     onmouseover="return searchBox.OnSearchSelectShow()"
     onmouseout="return searchBox.OnSearchSelectHide()"
     onkeydown="return searchBox.OnSearchSelectKey(event)">
</div>

<!-- iframe showing the search results (closed by default) -->
<div id="MSearchResultsWindow">
<iframe src="javascript:void(0)" frameborder="0" 
        name="MSearchResults" id="MSearchResults">
</iframe>
</div>

<div class="header">
  <div class="headertitle">
<div class="title">TensorContractionThreadPool.h</div>  </div>
</div><!--header-->
<div class="contents">
<div class="fragment"><div class="line"><a name="l00001"></a><span class="lineno">    1</span>&#160;<span class="comment">// This file is part of Eigen, a lightweight C++ template library</span></div>
<div class="line"><a name="l00002"></a><span class="lineno">    2</span>&#160;<span class="comment">// for linear algebra.</span></div>
<div class="line"><a name="l00003"></a><span class="lineno">    3</span>&#160;<span class="comment">//</span></div>
<div class="line"><a name="l00004"></a><span class="lineno">    4</span>&#160;<span class="comment">// Copyright (C) 2014 Benoit Steiner &lt;benoit.steiner.goog@gmail.com&gt;</span></div>
<div class="line"><a name="l00005"></a><span class="lineno">    5</span>&#160;<span class="comment">//</span></div>
<div class="line"><a name="l00006"></a><span class="lineno">    6</span>&#160;<span class="comment">// This Source Code Form is subject to the terms of the Mozilla</span></div>
<div class="line"><a name="l00007"></a><span class="lineno">    7</span>&#160;<span class="comment">// Public License v. 2.0. If a copy of the MPL was not distributed</span></div>
<div class="line"><a name="l00008"></a><span class="lineno">    8</span>&#160;<span class="comment">// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.</span></div>
<div class="line"><a name="l00009"></a><span class="lineno">    9</span>&#160; </div>
<div class="line"><a name="l00010"></a><span class="lineno">   10</span>&#160;<span class="preprocessor">#ifndef EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_THREAD_POOL_H</span></div>
<div class="line"><a name="l00011"></a><span class="lineno">   11</span>&#160;<span class="preprocessor">#define EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_THREAD_POOL_H</span></div>
<div class="line"><a name="l00012"></a><span class="lineno">   12</span>&#160; </div>
<div class="line"><a name="l00013"></a><span class="lineno">   13</span>&#160;<span class="comment">// evaluator for thread pool device</span></div>
<div class="line"><a name="l00014"></a><span class="lineno">   14</span>&#160;<span class="preprocessor">#ifdef EIGEN_USE_THREADS</span></div>
<div class="line"><a name="l00015"></a><span class="lineno">   15</span>&#160; </div>
<div class="line"><a name="l00016"></a><span class="lineno">   16</span>&#160;<span class="preprocessor">#include &quot;./InternalHeaderCheck.h&quot;</span></div>
<div class="line"><a name="l00017"></a><span class="lineno">   17</span>&#160; </div>
<div class="line"><a name="l00018"></a><span class="lineno">   18</span>&#160;<span class="keyword">namespace </span><a class="code" href="namespaceEigen.html">Eigen</a> {</div>
<div class="line"><a name="l00019"></a><span class="lineno">   19</span>&#160; </div>
<div class="line"><a name="l00020"></a><span class="lineno">   20</span>&#160;<span class="keyword">template</span>&lt;<span class="keyword">typename</span> Indices, <span class="keyword">typename</span> LeftArgType, <span class="keyword">typename</span> RightArgType, <span class="keyword">typename</span> OutputKernelType&gt;</div>
<div class="line"><a name="l00021"></a><span class="lineno">   21</span>&#160;<span class="keyword">struct </span>TensorEvaluator&lt;const TensorContractionOp&lt;Indices, LeftArgType, RightArgType, OutputKernelType&gt;, ThreadPoolDevice&gt; :</div>
<div class="line"><a name="l00022"></a><span class="lineno">   22</span>&#160;    <span class="keyword">public</span> TensorContractionEvaluatorBase&lt;TensorEvaluator&lt;const TensorContractionOp&lt;Indices, LeftArgType, RightArgType, OutputKernelType&gt;, ThreadPoolDevice&gt; &gt; {</div>
<div class="line"><a name="l00023"></a><span class="lineno">   23</span>&#160; </div>
<div class="line"><a name="l00024"></a><span class="lineno">   24</span>&#160;  <span class="keyword">typedef</span> ThreadPoolDevice Device;</div>
<div class="line"><a name="l00025"></a><span class="lineno">   25</span>&#160; </div>
<div class="line"><a name="l00026"></a><span class="lineno">   26</span>&#160;  <span class="keyword">typedef</span> TensorEvaluator&lt;const TensorContractionOp&lt;Indices, LeftArgType, RightArgType, OutputKernelType&gt;, Device&gt; Self;</div>
<div class="line"><a name="l00027"></a><span class="lineno">   27</span>&#160;  <span class="keyword">typedef</span> TensorContractionEvaluatorBase&lt;Self&gt; Base;</div>
<div class="line"><a name="l00028"></a><span class="lineno">   28</span>&#160; </div>
<div class="line"><a name="l00029"></a><span class="lineno">   29</span>&#160;  <span class="keyword">typedef</span> TensorContractionOp&lt;Indices, LeftArgType, RightArgType, OutputKernelType&gt; XprType;</div>
<div class="line"><a name="l00030"></a><span class="lineno">   30</span>&#160;  <span class="keyword">typedef</span> std::remove_const_t&lt;typename XprType::Scalar&gt; Scalar;</div>
<div class="line"><a name="l00031"></a><span class="lineno">   31</span>&#160;  <span class="keyword">typedef</span> <span class="keyword">typename</span> XprType::Index Index;</div>
<div class="line"><a name="l00032"></a><span class="lineno">   32</span>&#160;  <span class="keyword">typedef</span> <span class="keyword">typename</span> XprType::CoeffReturnType CoeffReturnType;</div>
<div class="line"><a name="l00033"></a><span class="lineno">   33</span>&#160;  <span class="keyword">typedef</span> <span class="keyword">typename</span> PacketType&lt;CoeffReturnType, Device&gt;::type PacketReturnType;</div>
<div class="line"><a name="l00034"></a><span class="lineno">   34</span>&#160; </div>
<div class="line"><a name="l00035"></a><span class="lineno">   35</span>&#160;  <span class="keyword">static</span> constexpr <span class="keywordtype">int</span> Layout = TensorEvaluator&lt;LeftArgType, Device&gt;::Layout;</div>
<div class="line"><a name="l00036"></a><span class="lineno">   36</span>&#160; </div>
<div class="line"><a name="l00037"></a><span class="lineno">   37</span>&#160;  <span class="comment">// Most of the code is assuming that both input tensors are ColMajor. If the</span></div>
<div class="line"><a name="l00038"></a><span class="lineno">   38</span>&#160;  <span class="comment">// inputs are RowMajor, we will &quot;cheat&quot; by swapping the LHS and RHS:</span></div>
<div class="line"><a name="l00039"></a><span class="lineno">   39</span>&#160;  <span class="comment">// If we want to compute A * B = C, where A is LHS and B is RHS, the code</span></div>
<div class="line"><a name="l00040"></a><span class="lineno">   40</span>&#160;  <span class="comment">// will pretend B is LHS and A is RHS.</span></div>
<div class="line"><a name="l00041"></a><span class="lineno">   41</span>&#160;  <span class="keyword">typedef</span> std::conditional_t&lt;</div>
<div class="line"><a name="l00042"></a><span class="lineno">   42</span>&#160;    <span class="keyword">static_cast&lt;</span><span class="keywordtype">int</span><span class="keyword">&gt;</span>(Layout) == <span class="keyword">static_cast&lt;</span><span class="keywordtype">int</span><span class="keyword">&gt;</span>(<a class="codeRef" href="../group__enums.html#ggaacded1a18ae58b0f554751f6cdf9eb13a0103672ae41005ab03b4176c765afd62">ColMajor</a>), LeftArgType, RightArgType&gt; EvalLeftArgType;</div>
<div class="line"><a name="l00043"></a><span class="lineno">   43</span>&#160;  <span class="keyword">typedef</span> std::conditional_t&lt;</div>
<div class="line"><a name="l00044"></a><span class="lineno">   44</span>&#160;    <span class="keyword">static_cast&lt;</span><span class="keywordtype">int</span><span class="keyword">&gt;</span>(Layout) == <span class="keyword">static_cast&lt;</span><span class="keywordtype">int</span><span class="keyword">&gt;</span>(<a class="codeRef" href="../group__enums.html#ggaacded1a18ae58b0f554751f6cdf9eb13a0103672ae41005ab03b4176c765afd62">ColMajor</a>), RightArgType, LeftArgType&gt; EvalRightArgType;</div>
<div class="line"><a name="l00045"></a><span class="lineno">   45</span>&#160; </div>
<div class="line"><a name="l00046"></a><span class="lineno">   46</span>&#160;  <span class="keyword">static</span> constexpr <span class="keywordtype">int</span> LDims =</div>
<div class="line"><a name="l00047"></a><span class="lineno">   47</span>&#160;      internal::array_size&lt;typename TensorEvaluator&lt;EvalLeftArgType, Device&gt;::Dimensions&gt;::value;</div>
<div class="line"><a name="l00048"></a><span class="lineno">   48</span>&#160;  <span class="keyword">static</span> constexpr <span class="keywordtype">int</span> RDims =</div>
<div class="line"><a name="l00049"></a><span class="lineno">   49</span>&#160;      internal::array_size&lt;typename TensorEvaluator&lt;EvalRightArgType, Device&gt;::Dimensions&gt;::value;</div>
<div class="line"><a name="l00050"></a><span class="lineno">   50</span>&#160;  <span class="keyword">static</span> constexpr <span class="keywordtype">int</span> ContractDims = internal::array_size&lt;Indices&gt;::value;</div>
<div class="line"><a name="l00051"></a><span class="lineno">   51</span>&#160; </div>
<div class="line"><a name="l00052"></a><span class="lineno">   52</span>&#160;  <span class="keyword">typedef</span> array&lt;Index, LDims&gt; left_dim_mapper_t;</div>
<div class="line"><a name="l00053"></a><span class="lineno">   53</span>&#160;  <span class="keyword">typedef</span> array&lt;Index, RDims&gt; right_dim_mapper_t;</div>
<div class="line"><a name="l00054"></a><span class="lineno">   54</span>&#160; </div>
<div class="line"><a name="l00055"></a><span class="lineno">   55</span>&#160;  <span class="keyword">typedef</span> array&lt;Index, ContractDims&gt; contract_t;</div>
<div class="line"><a name="l00056"></a><span class="lineno">   56</span>&#160;  <span class="keyword">typedef</span> array&lt;<a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>, LDims - ContractDims&gt; left_nocontract_t;</div>
<div class="line"><a name="l00057"></a><span class="lineno">   57</span>&#160;  <span class="keyword">typedef</span> array&lt;<a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>, RDims - ContractDims&gt; right_nocontract_t;</div>
<div class="line"><a name="l00058"></a><span class="lineno">   58</span>&#160; </div>
<div class="line"><a name="l00059"></a><span class="lineno">   59</span>&#160;  <span class="keyword">static</span> constexpr <span class="keywordtype">int</span> NumDims = LDims + RDims - 2 * ContractDims;</div>
<div class="line"><a name="l00060"></a><span class="lineno">   60</span>&#160; </div>
<div class="line"><a name="l00061"></a><span class="lineno">   61</span>&#160;  <span class="keyword">typedef</span> DSizes&lt;Index, NumDims&gt; Dimensions;</div>
<div class="line"><a name="l00062"></a><span class="lineno">   62</span>&#160; </div>
<div class="line"><a name="l00063"></a><span class="lineno">   63</span>&#160;  <span class="comment">// typedefs needed in evalTo</span></div>
<div class="line"><a name="l00064"></a><span class="lineno">   64</span>&#160;  <span class="keyword">typedef</span> std::remove_const_t&lt;typename EvalLeftArgType::Scalar&gt; LhsScalar;</div>
<div class="line"><a name="l00065"></a><span class="lineno">   65</span>&#160;  <span class="keyword">typedef</span> std::remove_const_t&lt;typename EvalRightArgType::Scalar&gt; RhsScalar;</div>
<div class="line"><a name="l00066"></a><span class="lineno">   66</span>&#160;  <span class="keyword">typedef</span> <span class="keyword">typename</span> internal::gebp_traits&lt;LhsScalar, RhsScalar&gt; Traits;</div>
<div class="line"><a name="l00067"></a><span class="lineno">   67</span>&#160; </div>
<div class="line"><a name="l00068"></a><span class="lineno">   68</span>&#160;  <span class="keyword">typedef</span> TensorEvaluator&lt;EvalLeftArgType, Device&gt; LeftEvaluator;</div>
<div class="line"><a name="l00069"></a><span class="lineno">   69</span>&#160;  <span class="keyword">typedef</span> TensorEvaluator&lt;EvalRightArgType, Device&gt; RightEvaluator;</div>
<div class="line"><a name="l00070"></a><span class="lineno">   70</span>&#160; </div>
<div class="line"><a name="l00071"></a><span class="lineno">   71</span>&#160;  TensorEvaluator(<span class="keyword">const</span> XprType&amp; op, <span class="keyword">const</span> Device&amp; device) :</div>
<div class="line"><a name="l00072"></a><span class="lineno">   72</span>&#160;      Base(op, device) {}</div>
<div class="line"><a name="l00073"></a><span class="lineno">   73</span>&#160; </div>
<div class="line"><a name="l00074"></a><span class="lineno">   74</span>&#160;  <span class="keyword">template</span> &lt;<span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l00075"></a><span class="lineno">   75</span>&#160;  <span class="keywordtype">void</span> evalProduct(Scalar* buffer)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l00076"></a><span class="lineno">   76</span>&#160;    evalProductImpl&lt;NoCallback, Alignment&gt;(buffer, NoCallback());</div>
<div class="line"><a name="l00077"></a><span class="lineno">   77</span>&#160;  }</div>
<div class="line"><a name="l00078"></a><span class="lineno">   78</span>&#160; </div>
<div class="line"><a name="l00079"></a><span class="lineno">   79</span>&#160;  <span class="keyword">template</span> &lt;<span class="keyword">typename</span> EvalToCallback, <span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l00080"></a><span class="lineno">   80</span>&#160;  <span class="keywordtype">void</span> evalProductAsync(Scalar* buffer, EvalToCallback done)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l00081"></a><span class="lineno">   81</span>&#160;    evalProductImpl&lt;EvalToCallback, Alignment&gt;(buffer, std::move(done));</div>
<div class="line"><a name="l00082"></a><span class="lineno">   82</span>&#160;  }</div>
<div class="line"><a name="l00083"></a><span class="lineno">   83</span>&#160; </div>
<div class="line"><a name="l00084"></a><span class="lineno">   84</span>&#160;  <span class="keyword">template</span> &lt;<span class="keyword">typename</span> DoneCallback, <span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l00085"></a><span class="lineno">   85</span>&#160;  <span class="keywordtype">void</span> evalProductImpl(Scalar* buffer, DoneCallback done)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l00086"></a><span class="lineno">   86</span>&#160;    <span class="comment">// This function computes a lot of heuristics in multiple steps, and it</span></div>
<div class="line"><a name="l00087"></a><span class="lineno">   87</span>&#160;    <span class="comment">// also has multiple exit points. To keep it sane, readable and all in one</span></div>
<div class="line"><a name="l00088"></a><span class="lineno">   88</span>&#160;    <span class="comment">// place, sync/async execution decision is made at runtime at the very end.</span></div>
<div class="line"><a name="l00089"></a><span class="lineno">   89</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00090"></a><span class="lineno">   90</span>&#160;    <span class="comment">// (1) In sync mode we allocate Context on the stack, submit computations</span></div>
<div class="line"><a name="l00091"></a><span class="lineno">   91</span>&#160;    <span class="comment">//     to the device thread pool, and block on a barrier until it is</span></div>
<div class="line"><a name="l00092"></a><span class="lineno">   92</span>&#160;    <span class="comment">//     completed.</span></div>
<div class="line"><a name="l00093"></a><span class="lineno">   93</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00094"></a><span class="lineno">   94</span>&#160;    <span class="comment">// (2) In async mode we allocate Context on the heap, and after all tasks</span></div>
<div class="line"><a name="l00095"></a><span class="lineno">   95</span>&#160;    <span class="comment">//     are finished, we call provided the done callback, and delete a</span></div>
<div class="line"><a name="l00096"></a><span class="lineno">   96</span>&#160;    <span class="comment">//     context from the heap.</span></div>
<div class="line"><a name="l00097"></a><span class="lineno">   97</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00098"></a><span class="lineno">   98</span>&#160;    <span class="comment">// (*) EvalParallelContext &amp; EvalShardedByInnerDimContext owns all the state</span></div>
<div class="line"><a name="l00099"></a><span class="lineno">   99</span>&#160;    <span class="comment">// and temporary buffers, required for executing the tensor contraction.</span></div>
<div class="line"><a name="l00100"></a><span class="lineno">  100</span>&#160;    <span class="comment">// They are responsible for cleaning it up after contraction is done.</span></div>
<div class="line"><a name="l00101"></a><span class="lineno">  101</span>&#160;    <span class="keyword">static</span> <span class="keyword">const</span> <span class="keywordtype">bool</span> IsEvalInSyncMode =</div>
<div class="line"><a name="l00102"></a><span class="lineno">  102</span>&#160;        std::is_same&lt;DoneCallback, NoCallback&gt;::value;</div>
<div class="line"><a name="l00103"></a><span class="lineno">  103</span>&#160; </div>
<div class="line"><a name="l00104"></a><span class="lineno">  104</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> m = this-&gt;m_i_size;</div>
<div class="line"><a name="l00105"></a><span class="lineno">  105</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> n = this-&gt;m_j_size;</div>
<div class="line"><a name="l00106"></a><span class="lineno">  106</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> k = this-&gt;m_k_size;</div>
<div class="line"><a name="l00107"></a><span class="lineno">  107</span>&#160;    <span class="keywordflow">if</span> (m == 0 || n == 0 || k == 0) <span class="keywordflow">return</span>;</div>
<div class="line"><a name="l00108"></a><span class="lineno">  108</span>&#160; </div>
<div class="line"><a name="l00109"></a><span class="lineno">  109</span>&#160;    <span class="comment">// Compute a set of algorithm parameters:</span></div>
<div class="line"><a name="l00110"></a><span class="lineno">  110</span>&#160;    <span class="comment">// - kernel block sizes (bm, bn, bk)</span></div>
<div class="line"><a name="l00111"></a><span class="lineno">  111</span>&#160;    <span class="comment">// - task grain sizes (number of kernels executed per task: gm, gn)</span></div>
<div class="line"><a name="l00112"></a><span class="lineno">  112</span>&#160;    <span class="comment">// - number of threads</span></div>
<div class="line"><a name="l00113"></a><span class="lineno">  113</span>&#160;    <span class="comment">// - sharding by row/column</span></div>
<div class="line"><a name="l00114"></a><span class="lineno">  114</span>&#160;    <span class="comment">// - parallel packing or first lhs then rhs</span></div>
<div class="line"><a name="l00115"></a><span class="lineno">  115</span>&#160;    <span class="comment">// and some derived parameters:</span></div>
<div class="line"><a name="l00116"></a><span class="lineno">  116</span>&#160;    <span class="comment">// - number of tasks (nm, nn, nk)</span></div>
<div class="line"><a name="l00117"></a><span class="lineno">  117</span>&#160;    <span class="comment">// - number of kernels (nm0, nn0)</span></div>
<div class="line"><a name="l00118"></a><span class="lineno">  118</span>&#160;    <span class="comment">// Unfortunately, all these parameters are tightly interdependent.</span></div>
<div class="line"><a name="l00119"></a><span class="lineno">  119</span>&#160;    <span class="comment">// So in some cases we first compute approximate values, then compute other</span></div>
<div class="line"><a name="l00120"></a><span class="lineno">  120</span>&#160;    <span class="comment">// values based on these approximations and then refine the approximations.</span></div>
<div class="line"><a name="l00121"></a><span class="lineno">  121</span>&#160; </div>
<div class="line"><a name="l00122"></a><span class="lineno">  122</span>&#160;    <span class="comment">// There are lots of heuristics here. There is some reasoning behind them,</span></div>
<div class="line"><a name="l00123"></a><span class="lineno">  123</span>&#160;    <span class="comment">// but ultimately they are just tuned on contraction benchmarks for</span></div>
<div class="line"><a name="l00124"></a><span class="lineno">  124</span>&#160;    <span class="comment">// different input configurations, thread counts and instruction sets.</span></div>
<div class="line"><a name="l00125"></a><span class="lineno">  125</span>&#160;    <span class="comment">// So feel free to question any of them.</span></div>
<div class="line"><a name="l00126"></a><span class="lineno">  126</span>&#160; </div>
<div class="line"><a name="l00127"></a><span class="lineno">  127</span>&#160;    <span class="comment">// Compute whether we want to shard by row or by column.</span></div>
<div class="line"><a name="l00128"></a><span class="lineno">  128</span>&#160;    <span class="comment">// This is a first approximation, it will be refined later. Since we don&#39;t</span></div>
<div class="line"><a name="l00129"></a><span class="lineno">  129</span>&#160;    <span class="comment">// know number of threads yet we use 2, because what&#39;s we are most</span></div>
<div class="line"><a name="l00130"></a><span class="lineno">  130</span>&#160;    <span class="comment">// interested in at this point is whether it makes sense to use</span></div>
<div class="line"><a name="l00131"></a><span class="lineno">  131</span>&#160;    <span class="comment">// parallelization at all or not.</span></div>
<div class="line"><a name="l00132"></a><span class="lineno">  132</span>&#160;    <span class="keywordtype">bool</span> shard_by_col = shardByCol(m, n, 2);</div>
<div class="line"><a name="l00133"></a><span class="lineno">  133</span>&#160; </div>
<div class="line"><a name="l00134"></a><span class="lineno">  134</span>&#160;    <span class="comment">// First approximation of kernel blocking sizes.</span></div>
<div class="line"><a name="l00135"></a><span class="lineno">  135</span>&#160;    <span class="comment">// Again, we don&#39;t know number of threads yet, so we use 2.</span></div>
<div class="line"><a name="l00136"></a><span class="lineno">  136</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> bm, bn, bk;</div>
<div class="line"><a name="l00137"></a><span class="lineno">  137</span>&#160;    <span class="keywordflow">if</span> (shard_by_col) {</div>
<div class="line"><a name="l00138"></a><span class="lineno">  138</span>&#160;      internal::TensorContractionBlocking&lt;Scalar, LhsScalar, RhsScalar, <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>,</div>
<div class="line"><a name="l00139"></a><span class="lineno">  139</span>&#160;                                          internal::ShardByCol&gt;</div>
<div class="line"><a name="l00140"></a><span class="lineno">  140</span>&#160;          blocking(k, m, n, 2);</div>
<div class="line"><a name="l00141"></a><span class="lineno">  141</span>&#160;      bm = blocking.mc();</div>
<div class="line"><a name="l00142"></a><span class="lineno">  142</span>&#160;      bn = blocking.nc();</div>
<div class="line"><a name="l00143"></a><span class="lineno">  143</span>&#160;      bk = blocking.kc();</div>
<div class="line"><a name="l00144"></a><span class="lineno">  144</span>&#160;    } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00145"></a><span class="lineno">  145</span>&#160;      internal::TensorContractionBlocking&lt;Scalar, LhsScalar, RhsScalar, <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>,</div>
<div class="line"><a name="l00146"></a><span class="lineno">  146</span>&#160;                                          internal::ShardByRow&gt;</div>
<div class="line"><a name="l00147"></a><span class="lineno">  147</span>&#160;          blocking(k, m, n, 2);</div>
<div class="line"><a name="l00148"></a><span class="lineno">  148</span>&#160;      bm = blocking.mc();</div>
<div class="line"><a name="l00149"></a><span class="lineno">  149</span>&#160;      bn = blocking.nc();</div>
<div class="line"><a name="l00150"></a><span class="lineno">  150</span>&#160;      bk = blocking.kc();</div>
<div class="line"><a name="l00151"></a><span class="lineno">  151</span>&#160;    }</div>
<div class="line"><a name="l00152"></a><span class="lineno">  152</span>&#160; </div>
<div class="line"><a name="l00153"></a><span class="lineno">  153</span>&#160;    <span class="comment">// Compute optimal number of threads.</span></div>
<div class="line"><a name="l00154"></a><span class="lineno">  154</span>&#160;    <span class="comment">// Note: we use bk instead of k here because we are interested in amount of</span></div>
<div class="line"><a name="l00155"></a><span class="lineno">  155</span>&#160;    <span class="comment">// _parallelizable_ computations, and computations are not parallelizable</span></div>
<div class="line"><a name="l00156"></a><span class="lineno">  156</span>&#160;    <span class="comment">// across k dimension.</span></div>
<div class="line"><a name="l00157"></a><span class="lineno">  157</span>&#160;    <span class="keyword">const</span> TensorOpCost cost =</div>
<div class="line"><a name="l00158"></a><span class="lineno">  158</span>&#160;        contractionCost(m, n, bm, bn, bk, shard_by_col, <span class="keyword">false</span>);</div>
<div class="line"><a name="l00159"></a><span class="lineno">  159</span>&#160;    <span class="keywordtype">int</span> num_threads = TensorCostModel&lt;ThreadPoolDevice&gt;::numThreads(</div>
<div class="line"><a name="l00160"></a><span class="lineno">  160</span>&#160;        <span class="keyword">static_cast&lt;</span><span class="keywordtype">double</span><span class="keyword">&gt;</span>(n) * m, cost, this-&gt;m_device.numThreads());</div>
<div class="line"><a name="l00161"></a><span class="lineno">  161</span>&#160;    <span class="keywordtype">int</span> num_threads_by_k = numThreadsInnerDim(m, n, k);</div>
<div class="line"><a name="l00162"></a><span class="lineno">  162</span>&#160;    <span class="keywordflow">if</span> (shardByInnerDim(m, n, k, num_threads, num_threads_by_k)) {</div>
<div class="line"><a name="l00163"></a><span class="lineno">  163</span>&#160;      <span class="comment">// We are in the scenario where it is more effective to shard by the</span></div>
<div class="line"><a name="l00164"></a><span class="lineno">  164</span>&#160;      <span class="comment">// inner dimension.</span></div>
<div class="line"><a name="l00165"></a><span class="lineno">  165</span>&#160;      <span class="keywordflow">if</span> (IsEvalInSyncMode) {</div>
<div class="line"><a name="l00166"></a><span class="lineno">  166</span>&#160;        EvalShardedByInnerDimContext&lt;DoneCallback&gt; ctx(</div>
<div class="line"><a name="l00167"></a><span class="lineno">  167</span>&#160;            <span class="keyword">this</span>, num_threads_by_k, buffer, m, n, k, std::move(done));</div>
<div class="line"><a name="l00168"></a><span class="lineno">  168</span>&#160;        ctx.template run&lt;Alignment&gt;();</div>
<div class="line"><a name="l00169"></a><span class="lineno">  169</span>&#160;      } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00170"></a><span class="lineno">  170</span>&#160;        <span class="keyword">auto</span>* ctx = <span class="keyword">new</span> EvalShardedByInnerDimContext&lt;DoneCallback&gt;(</div>
<div class="line"><a name="l00171"></a><span class="lineno">  171</span>&#160;            <span class="keyword">this</span>, num_threads_by_k, buffer, m, n, k, std::move(done));</div>
<div class="line"><a name="l00172"></a><span class="lineno">  172</span>&#160;        ctx-&gt;template runAsync&lt;Alignment&gt;();</div>
<div class="line"><a name="l00173"></a><span class="lineno">  173</span>&#160;      }</div>
<div class="line"><a name="l00174"></a><span class="lineno">  174</span>&#160; </div>
<div class="line"><a name="l00175"></a><span class="lineno">  175</span>&#160;      <span class="keywordflow">return</span>;</div>
<div class="line"><a name="l00176"></a><span class="lineno">  176</span>&#160;    }</div>
<div class="line"><a name="l00177"></a><span class="lineno">  177</span>&#160; </div>
<div class="line"><a name="l00178"></a><span class="lineno">  178</span>&#160;    <span class="comment">// TODO(dvyukov): this is a stop-gap to prevent regressions while the cost</span></div>
<div class="line"><a name="l00179"></a><span class="lineno">  179</span>&#160;    <span class="comment">// model is not tuned. Remove this when the cost model is tuned.</span></div>
<div class="line"><a name="l00180"></a><span class="lineno">  180</span>&#160;    <span class="keywordflow">if</span> (n == 1) num_threads = 1;</div>
<div class="line"><a name="l00181"></a><span class="lineno">  181</span>&#160; </div>
<div class="line"><a name="l00182"></a><span class="lineno">  182</span>&#160;    <span class="keywordflow">if</span> (num_threads == 1) {</div>
<div class="line"><a name="l00183"></a><span class="lineno">  183</span>&#160;      TENSOR_CONTRACTION_DISPATCH(this-&gt;<span class="keyword">template</span> evalProductSequential,</div>
<div class="line"><a name="l00184"></a><span class="lineno">  184</span>&#160;                                  <a class="codeRef" href="../group__enums.html#gga45fe06e29902b7a2773de05ba27b47a1a4e19dd09d5ff42295ba1d72d12a46686">Unaligned</a>, (buffer));</div>
<div class="line"><a name="l00185"></a><span class="lineno">  185</span>&#160;      <span class="keywordflow">if</span> (!IsEvalInSyncMode) done();</div>
<div class="line"><a name="l00186"></a><span class="lineno">  186</span>&#160;      <span class="keywordflow">return</span>;</div>
<div class="line"><a name="l00187"></a><span class="lineno">  187</span>&#160;    }</div>
<div class="line"><a name="l00188"></a><span class="lineno">  188</span>&#160; </div>
<div class="line"><a name="l00189"></a><span class="lineno">  189</span>&#160;    <span class="comment">// Now that we know number of threads, recalculate sharding and blocking.</span></div>
<div class="line"><a name="l00190"></a><span class="lineno">  190</span>&#160;    shard_by_col = shardByCol(m, n, num_threads);</div>
<div class="line"><a name="l00191"></a><span class="lineno">  191</span>&#160;    <span class="keywordflow">if</span> (shard_by_col) {</div>
<div class="line"><a name="l00192"></a><span class="lineno">  192</span>&#160;      internal::TensorContractionBlocking&lt;Scalar, LhsScalar, RhsScalar, <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>,</div>
<div class="line"><a name="l00193"></a><span class="lineno">  193</span>&#160;                                          internal::ShardByCol&gt;</div>
<div class="line"><a name="l00194"></a><span class="lineno">  194</span>&#160;          blocking(k, m, n, num_threads);</div>
<div class="line"><a name="l00195"></a><span class="lineno">  195</span>&#160;      bm = blocking.mc();</div>
<div class="line"><a name="l00196"></a><span class="lineno">  196</span>&#160;      bn = blocking.nc();</div>
<div class="line"><a name="l00197"></a><span class="lineno">  197</span>&#160;      bk = blocking.kc();</div>
<div class="line"><a name="l00198"></a><span class="lineno">  198</span>&#160;    } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00199"></a><span class="lineno">  199</span>&#160;      internal::TensorContractionBlocking&lt;Scalar, LhsScalar, RhsScalar, <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>,</div>
<div class="line"><a name="l00200"></a><span class="lineno">  200</span>&#160;                                          internal::ShardByRow&gt;</div>
<div class="line"><a name="l00201"></a><span class="lineno">  201</span>&#160;          blocking(k, m, n, num_threads);</div>
<div class="line"><a name="l00202"></a><span class="lineno">  202</span>&#160;      bm = blocking.mc();</div>
<div class="line"><a name="l00203"></a><span class="lineno">  203</span>&#160;      bn = blocking.nc();</div>
<div class="line"><a name="l00204"></a><span class="lineno">  204</span>&#160;      bk = blocking.kc();</div>
<div class="line"><a name="l00205"></a><span class="lineno">  205</span>&#160;    }</div>
<div class="line"><a name="l00206"></a><span class="lineno">  206</span>&#160; </div>
<div class="line"><a name="l00207"></a><span class="lineno">  207</span>&#160;    <span class="comment">// Number of kernels for each dimension.</span></div>
<div class="line"><a name="l00208"></a><span class="lineno">  208</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nm0 = divup(m, bm);</div>
<div class="line"><a name="l00209"></a><span class="lineno">  209</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nn0 = divup(n, bn);</div>
<div class="line"><a name="l00210"></a><span class="lineno">  210</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nk = divup(k, bk);</div>
<div class="line"><a name="l00211"></a><span class="lineno">  211</span>&#160; </div>
<div class="line"><a name="l00212"></a><span class="lineno">  212</span>&#160;    <span class="comment">// Calculate task grain size (number of kernels executed per task).</span></div>
<div class="line"><a name="l00213"></a><span class="lineno">  213</span>&#160;    <span class="comment">// This task size coarsening serves two purposes:</span></div>
<div class="line"><a name="l00214"></a><span class="lineno">  214</span>&#160;    <span class="comment">// 1. It reduces per-task overheads including synchronization overheads.</span></div>
<div class="line"><a name="l00215"></a><span class="lineno">  215</span>&#160;    <span class="comment">// 2. It allows to use caches better (reuse the same packed rhs in several</span></div>
<div class="line"><a name="l00216"></a><span class="lineno">  216</span>&#160;    <span class="comment">// consecutive kernels).</span></div>
<div class="line"><a name="l00217"></a><span class="lineno">  217</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> gm = 1;</div>
<div class="line"><a name="l00218"></a><span class="lineno">  218</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> gn = 1;</div>
<div class="line"><a name="l00219"></a><span class="lineno">  219</span>&#160;    <span class="comment">// If we are sharding by column, then we prefer to reduce rows first.</span></div>
<div class="line"><a name="l00220"></a><span class="lineno">  220</span>&#160;    <span class="keywordflow">if</span> (shard_by_col) {</div>
<div class="line"><a name="l00221"></a><span class="lineno">  221</span>&#160;      gm = coarsenM(m, n, bm, bn, bk, gn, num_threads, shard_by_col);</div>
<div class="line"><a name="l00222"></a><span class="lineno">  222</span>&#160;      gn = coarsenN(m, n, bm, bn, bk, gm, num_threads, shard_by_col);</div>
<div class="line"><a name="l00223"></a><span class="lineno">  223</span>&#160;    } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00224"></a><span class="lineno">  224</span>&#160;      gn = coarsenN(m, n, bm, bn, bk, gm, num_threads, shard_by_col);</div>
<div class="line"><a name="l00225"></a><span class="lineno">  225</span>&#160;      gm = coarsenM(m, n, bm, bn, bk, gn, num_threads, shard_by_col);</div>
<div class="line"><a name="l00226"></a><span class="lineno">  226</span>&#160;    }</div>
<div class="line"><a name="l00227"></a><span class="lineno">  227</span>&#160;    <span class="comment">// Number of tasks in each dimension.</span></div>
<div class="line"><a name="l00228"></a><span class="lineno">  228</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nm = divup(nm0, gm);</div>
<div class="line"><a name="l00229"></a><span class="lineno">  229</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nn = divup(nn0, gn);</div>
<div class="line"><a name="l00230"></a><span class="lineno">  230</span>&#160; </div>
<div class="line"><a name="l00231"></a><span class="lineno">  231</span>&#160;    <span class="comment">// If there is enough concurrency in the sharding dimension, we choose not</span></div>
<div class="line"><a name="l00232"></a><span class="lineno">  232</span>&#160;    <span class="comment">// to paralellize by the other dimension, and execute all kernels in sync</span></div>
<div class="line"><a name="l00233"></a><span class="lineno">  233</span>&#160;    <span class="comment">// mode. This reduces parallelism from the nm x nn down to nn</span></div>
<div class="line"><a name="l00234"></a><span class="lineno">  234</span>&#160;    <span class="comment">// (shard_by_col==true) or nm (shard_by_col==false).</span></div>
<div class="line"><a name="l00235"></a><span class="lineno">  235</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> sharding_dim_tasks = shard_by_col ? nn : nm;</div>
<div class="line"><a name="l00236"></a><span class="lineno">  236</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">int</span> num_worker_threads = this-&gt;m_device.numThreadsInPool();</div>
<div class="line"><a name="l00237"></a><span class="lineno">  237</span>&#160; </div>
<div class="line"><a name="l00238"></a><span class="lineno">  238</span>&#160;    <span class="comment">// With small number of threads we want to make sure that we do not reduce</span></div>
<div class="line"><a name="l00239"></a><span class="lineno">  239</span>&#160;    <span class="comment">// parallelism too much. With large number of threads we trade maximum</span></div>
<div class="line"><a name="l00240"></a><span class="lineno">  240</span>&#160;    <span class="comment">// parallelism for better memory locality.</span></div>
<div class="line"><a name="l00241"></a><span class="lineno">  241</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">float</span> oversharding_factor =</div>
<div class="line"><a name="l00242"></a><span class="lineno">  242</span>&#160;        num_worker_threads &lt;= 4  ? 8.0 :</div>
<div class="line"><a name="l00243"></a><span class="lineno">  243</span>&#160;        num_worker_threads &lt;= 8  ? 4.0 :</div>
<div class="line"><a name="l00244"></a><span class="lineno">  244</span>&#160;        num_worker_threads &lt;= 16 ? 2.0 :</div>
<div class="line"><a name="l00245"></a><span class="lineno">  245</span>&#160;        num_worker_threads &lt;= 32 ? 1.0 :</div>
<div class="line"><a name="l00246"></a><span class="lineno">  246</span>&#160;        num_worker_threads &lt;= 64 ? 0.8 : <span class="comment">/* num_worker_threads &gt; 64 */</span> 0.6;</div>
<div class="line"><a name="l00247"></a><span class="lineno">  247</span>&#160; </div>
<div class="line"><a name="l00248"></a><span class="lineno">  248</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">bool</span> parallelize_by_sharding_dim_only =</div>
<div class="line"><a name="l00249"></a><span class="lineno">  249</span>&#160;        sharding_dim_tasks &gt;= oversharding_factor * num_worker_threads;</div>
<div class="line"><a name="l00250"></a><span class="lineno">  250</span>&#160; </div>
<div class="line"><a name="l00251"></a><span class="lineno">  251</span>&#160;    <span class="comment">// Last by not least, decide whether we want to issue both lhs and rhs</span></div>
<div class="line"><a name="l00252"></a><span class="lineno">  252</span>&#160;    <span class="comment">// packing in parallel; or issue lhs packing first, and then issue rhs</span></div>
<div class="line"><a name="l00253"></a><span class="lineno">  253</span>&#160;    <span class="comment">// packing when lhs packing completes (for !shard_by_col lhs and rhs are</span></div>
<div class="line"><a name="l00254"></a><span class="lineno">  254</span>&#160;    <span class="comment">// swapped). Parallel packing allows more parallelism (for both packing and</span></div>
<div class="line"><a name="l00255"></a><span class="lineno">  255</span>&#160;    <span class="comment">// kernels), while sequential packing provides better locality (once</span></div>
<div class="line"><a name="l00256"></a><span class="lineno">  256</span>&#160;    <span class="comment">// a thread finishes rhs packing it proceed to kernels with that rhs).</span></div>
<div class="line"><a name="l00257"></a><span class="lineno">  257</span>&#160;    <span class="comment">// First, we are interested in parallel packing if there are few tasks.</span></div>
<div class="line"><a name="l00258"></a><span class="lineno">  258</span>&#160;    <span class="keywordtype">bool</span> parallel_pack = num_threads &gt;= nm * nn;</div>
<div class="line"><a name="l00259"></a><span class="lineno">  259</span>&#160;    <span class="comment">// Also do parallel packing if all data fits into L2$.</span></div>
<div class="line"><a name="l00260"></a><span class="lineno">  260</span>&#160;    <span class="keywordflow">if</span> (m * bk * <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>(<span class="keyword">sizeof</span>(LhsScalar)) + n * bk * <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>(<span class="keyword">sizeof</span>(RhsScalar)) &lt;=</div>
<div class="line"><a name="l00261"></a><span class="lineno">  261</span>&#160;        l2CacheSize() * num_threads)</div>
<div class="line"><a name="l00262"></a><span class="lineno">  262</span>&#160;      parallel_pack = <span class="keyword">true</span>;</div>
<div class="line"><a name="l00263"></a><span class="lineno">  263</span>&#160;    <span class="comment">// But don&#39;t do it if we will use each rhs only once. Locality seems to be</span></div>
<div class="line"><a name="l00264"></a><span class="lineno">  264</span>&#160;    <span class="comment">// more important in this case.</span></div>
<div class="line"><a name="l00265"></a><span class="lineno">  265</span>&#160;    <span class="keywordflow">if</span> ((shard_by_col ? nm : nn) == 1) parallel_pack = <span class="keyword">false</span>;</div>
<div class="line"><a name="l00266"></a><span class="lineno">  266</span>&#160;    <span class="comment">// Also don&#39;t get in the way of parallelize_by_sharding_dim_only</span></div>
<div class="line"><a name="l00267"></a><span class="lineno">  267</span>&#160;    <span class="comment">// optimization.</span></div>
<div class="line"><a name="l00268"></a><span class="lineno">  268</span>&#160;    <span class="keywordflow">if</span> (parallelize_by_sharding_dim_only) parallel_pack = <span class="keyword">false</span>;</div>
<div class="line"><a name="l00269"></a><span class="lineno">  269</span>&#160; </div>
<div class="line"><a name="l00270"></a><span class="lineno">  270</span>&#160;    <span class="comment">// TODO(ezhulnev): With if contexpr we don&#39;t need SyncEvalParallelContext.</span></div>
<div class="line"><a name="l00271"></a><span class="lineno">  271</span>&#160;    <span class="keywordflow">if</span> (IsEvalInSyncMode) {</div>
<div class="line"><a name="l00272"></a><span class="lineno">  272</span>&#160;<span class="preprocessor">#define CONTEXT_ARGS                                                        \</span></div>
<div class="line"><a name="l00273"></a><span class="lineno">  273</span>&#160;<span class="preprocessor">  (this, num_threads, buffer, m, n, k, bm, bn, bk, nm, nn, nk, gm, gn, nm0, \</span></div>
<div class="line"><a name="l00274"></a><span class="lineno">  274</span>&#160;<span class="preprocessor">   nn0, shard_by_col, parallel_pack, parallelize_by_sharding_dim_only,      \</span></div>
<div class="line"><a name="l00275"></a><span class="lineno">  275</span>&#160;<span class="preprocessor">   NoCallback())                                                            \</span></div>
<div class="line"><a name="l00276"></a><span class="lineno">  276</span>&#160;<span class="preprocessor">      .run()</span></div>
<div class="line"><a name="l00277"></a><span class="lineno">  277</span>&#160;      TENSOR_CONTRACTION_DISPATCH(SyncEvalParallelContext, Alignment,</div>
<div class="line"><a name="l00278"></a><span class="lineno">  278</span>&#160;                                  CONTEXT_ARGS);</div>
<div class="line"><a name="l00279"></a><span class="lineno">  279</span>&#160;<span class="preprocessor">#undef CONTEXT_ARGS</span></div>
<div class="line"><a name="l00280"></a><span class="lineno">  280</span>&#160; </div>
<div class="line"><a name="l00281"></a><span class="lineno">  281</span>&#160;    } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00282"></a><span class="lineno">  282</span>&#160;<span class="preprocessor">#define CONTEXT_ARGS                                                        \</span></div>
<div class="line"><a name="l00283"></a><span class="lineno">  283</span>&#160;<span class="preprocessor">  (this, num_threads, buffer, m, n, k, bm, bn, bk, nm, nn, nk, gm, gn, nm0, \</span></div>
<div class="line"><a name="l00284"></a><span class="lineno">  284</span>&#160;<span class="preprocessor">   nn0, shard_by_col, parallel_pack, parallelize_by_sharding_dim_only,      \</span></div>
<div class="line"><a name="l00285"></a><span class="lineno">  285</span>&#160;<span class="preprocessor">   std::move(done))</span></div>
<div class="line"><a name="l00286"></a><span class="lineno">  286</span>&#160;      TENSOR_CONTRACTION_ASYNC_DISPATCH(EvalParallelContext, DoneCallback,</div>
<div class="line"><a name="l00287"></a><span class="lineno">  287</span>&#160;                                        Alignment, CONTEXT_ARGS, run());</div>
<div class="line"><a name="l00288"></a><span class="lineno">  288</span>&#160;<span class="preprocessor">#undef CONTEXT_ARGS</span></div>
<div class="line"><a name="l00289"></a><span class="lineno">  289</span>&#160;    }</div>
<div class="line"><a name="l00290"></a><span class="lineno">  290</span>&#160;  }</div>
<div class="line"><a name="l00291"></a><span class="lineno">  291</span>&#160; </div>
<div class="line"><a name="l00292"></a><span class="lineno">  292</span>&#160;  <span class="comment">// ------------------------------------------------------------------------ //</span></div>
<div class="line"><a name="l00293"></a><span class="lineno">  293</span>&#160; </div>
<div class="line"><a name="l00294"></a><span class="lineno">  294</span>&#160;  <span class="comment">// Dummy struct to represent an empty DoneCallback.</span></div>
<div class="line"><a name="l00295"></a><span class="lineno">  295</span>&#160; </div>
<div class="line"><a name="l00296"></a><span class="lineno">  296</span>&#160;  <span class="keyword">struct </span>NoCallback {</div>
<div class="line"><a name="l00297"></a><span class="lineno">  297</span>&#160;    <span class="keywordtype">void</span> operator()() {</div>
<div class="line"><a name="l00298"></a><span class="lineno">  298</span>&#160;      eigen_assert(<span class="keyword">false</span> &amp;&amp; <span class="stringliteral">&quot;NoCallback should never be called&quot;</span>);</div>
<div class="line"><a name="l00299"></a><span class="lineno">  299</span>&#160;    }</div>
<div class="line"><a name="l00300"></a><span class="lineno">  300</span>&#160;  };</div>
<div class="line"><a name="l00301"></a><span class="lineno">  301</span>&#160; </div>
<div class="line"><a name="l00302"></a><span class="lineno">  302</span>&#160;  <span class="comment">// ------------------------------------------------------------------------ //</span></div>
<div class="line"><a name="l00303"></a><span class="lineno">  303</span>&#160; </div>
<div class="line"><a name="l00304"></a><span class="lineno">  304</span>&#160;  <span class="keyword">template</span> &lt;<span class="keyword">typename</span> DoneCallback, <span class="keyword">typename</span> Context&gt;</div>
<div class="line"><a name="l00305"></a><span class="lineno">  305</span>&#160;  <span class="keyword">class </span>EvalParallelNotification;</div>
<div class="line"><a name="l00306"></a><span class="lineno">  306</span>&#160; </div>
<div class="line"><a name="l00307"></a><span class="lineno">  307</span>&#160;  <span class="comment">// Synchronous evaluation notification that blocks caller thread in Wait().</span></div>
<div class="line"><a name="l00308"></a><span class="lineno">  308</span>&#160;  <span class="keyword">template</span> &lt;<span class="keyword">typename</span> Context&gt;</div>
<div class="line"><a name="l00309"></a><span class="lineno">  309</span>&#160;  <span class="keyword">class </span>EvalParallelNotification&lt;NoCallback, Context&gt; {</div>
<div class="line"><a name="l00310"></a><span class="lineno">  310</span>&#160;   <span class="keyword">public</span>:</div>
<div class="line"><a name="l00311"></a><span class="lineno">  311</span>&#160;    EvalParallelNotification(Context*, NoCallback) {}</div>
<div class="line"><a name="l00312"></a><span class="lineno">  312</span>&#160;    <span class="keywordtype">void</span> Notify() { done_.Notify(); }</div>
<div class="line"><a name="l00313"></a><span class="lineno">  313</span>&#160;    <span class="keywordtype">void</span> Wait() { done_.Wait(); }</div>
<div class="line"><a name="l00314"></a><span class="lineno">  314</span>&#160;   <span class="keyword">private</span>:</div>
<div class="line"><a name="l00315"></a><span class="lineno">  315</span>&#160;    Eigen::Notification done_;</div>
<div class="line"><a name="l00316"></a><span class="lineno">  316</span>&#160;  };</div>
<div class="line"><a name="l00317"></a><span class="lineno">  317</span>&#160; </div>
<div class="line"><a name="l00318"></a><span class="lineno">  318</span>&#160;  <span class="comment">// Asynchronous evaluation notification that does not block in Wait().</span></div>
<div class="line"><a name="l00319"></a><span class="lineno">  319</span>&#160;  <span class="keyword">template</span> &lt;<span class="keyword">typename</span> DoneCallback, <span class="keyword">typename</span> Context&gt;</div>
<div class="line"><a name="l00320"></a><span class="lineno">  320</span>&#160;  <span class="keyword">class </span>EvalParallelNotification {</div>
<div class="line"><a name="l00321"></a><span class="lineno">  321</span>&#160;   <span class="keyword">public</span>:</div>
<div class="line"><a name="l00322"></a><span class="lineno">  322</span>&#160;    EvalParallelNotification(Context* ctx, DoneCallback done)</div>
<div class="line"><a name="l00323"></a><span class="lineno">  323</span>&#160;        : ctx_(ctx), done_(std::move(done)) {}</div>
<div class="line"><a name="l00324"></a><span class="lineno">  324</span>&#160; </div>
<div class="line"><a name="l00325"></a><span class="lineno">  325</span>&#160;    <span class="keywordtype">void</span> Notify() {</div>
<div class="line"><a name="l00326"></a><span class="lineno">  326</span>&#160;      <span class="comment">// Make a copy of done callback, because it will be destructed when we</span></div>
<div class="line"><a name="l00327"></a><span class="lineno">  327</span>&#160;      <span class="comment">// will delete context in the next line (EvalParallelNotification is a</span></div>
<div class="line"><a name="l00328"></a><span class="lineno">  328</span>&#160;      <span class="comment">// data member of EvalParallelContext class).</span></div>
<div class="line"><a name="l00329"></a><span class="lineno">  329</span>&#160;      DoneCallback done_copy = std::move(done_);</div>
<div class="line"><a name="l00330"></a><span class="lineno">  330</span>&#160; </div>
<div class="line"><a name="l00331"></a><span class="lineno">  331</span>&#160;      <span class="comment">// Delete parallel evaluation context.</span></div>
<div class="line"><a name="l00332"></a><span class="lineno">  332</span>&#160;      <span class="keyword">delete</span> ctx_;</div>
<div class="line"><a name="l00333"></a><span class="lineno">  333</span>&#160; </div>
<div class="line"><a name="l00334"></a><span class="lineno">  334</span>&#160;      <span class="comment">// Now safely call the done callback.</span></div>
<div class="line"><a name="l00335"></a><span class="lineno">  335</span>&#160;      done_copy();</div>
<div class="line"><a name="l00336"></a><span class="lineno">  336</span>&#160;    }</div>
<div class="line"><a name="l00337"></a><span class="lineno">  337</span>&#160; </div>
<div class="line"><a name="l00338"></a><span class="lineno">  338</span>&#160;    <span class="keywordtype">void</span> Wait() {}</div>
<div class="line"><a name="l00339"></a><span class="lineno">  339</span>&#160; </div>
<div class="line"><a name="l00340"></a><span class="lineno">  340</span>&#160;   <span class="keyword">private</span>:</div>
<div class="line"><a name="l00341"></a><span class="lineno">  341</span>&#160;    Context* ctx_;</div>
<div class="line"><a name="l00342"></a><span class="lineno">  342</span>&#160;    DoneCallback done_;</div>
<div class="line"><a name="l00343"></a><span class="lineno">  343</span>&#160;  };</div>
<div class="line"><a name="l00344"></a><span class="lineno">  344</span>&#160; </div>
<div class="line"><a name="l00345"></a><span class="lineno">  345</span>&#160;  <span class="comment">// Context orchestrates sync/async parallel contraction evaluation. When it is</span></div>
<div class="line"><a name="l00346"></a><span class="lineno">  346</span>&#160;  <span class="comment">// executed in asynchronous mode, it owns all the shared state that might be</span></div>
<div class="line"><a name="l00347"></a><span class="lineno">  347</span>&#160;  <span class="comment">// accessible by block packing and kernel tasks.</span></div>
<div class="line"><a name="l00348"></a><span class="lineno">  348</span>&#160; </div>
<div class="line"><a name="l00349"></a><span class="lineno">  349</span>&#160;  <span class="keyword">template</span> &lt;<span class="keyword">typename</span> DoneCallback, <span class="keywordtype">bool</span> lhs_inner_dim_contiguous,</div>
<div class="line"><a name="l00350"></a><span class="lineno">  350</span>&#160;            <span class="keywordtype">bool</span> rhs_inner_dim_contiguous, <span class="keywordtype">bool</span> rhs_inner_dim_reordered,</div>
<div class="line"><a name="l00351"></a><span class="lineno">  351</span>&#160;            <span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l00352"></a><span class="lineno">  352</span>&#160;  <span class="keyword">class </span>EvalParallelContext {</div>
<div class="line"><a name="l00353"></a><span class="lineno">  353</span>&#160;   <span class="keyword">public</span>:</div>
<div class="line"><a name="l00354"></a><span class="lineno">  354</span>&#160;    <span class="keyword">typedef</span> internal::TensorContractionInputMapper&lt;</div>
<div class="line"><a name="l00355"></a><span class="lineno">  355</span>&#160;        LhsScalar, <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>, internal::Lhs, LeftEvaluator, left_nocontract_t,</div>
<div class="line"><a name="l00356"></a><span class="lineno">  356</span>&#160;        contract_t, internal::packet_traits&lt;LhsScalar&gt;::size,</div>
<div class="line"><a name="l00357"></a><span class="lineno">  357</span>&#160;        lhs_inner_dim_contiguous, <span class="keyword">false</span>, <a class="codeRef" href="../group__enums.html#gga45fe06e29902b7a2773de05ba27b47a1a4e19dd09d5ff42295ba1d72d12a46686">Unaligned</a>&gt;</div>
<div class="line"><a name="l00358"></a><span class="lineno">  358</span>&#160;        LhsMapper;</div>
<div class="line"><a name="l00359"></a><span class="lineno">  359</span>&#160;    <span class="keyword">typedef</span> internal::TensorContractionInputMapper&lt;</div>
<div class="line"><a name="l00360"></a><span class="lineno">  360</span>&#160;        RhsScalar, <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>, internal::Rhs, RightEvaluator, right_nocontract_t,</div>
<div class="line"><a name="l00361"></a><span class="lineno">  361</span>&#160;        contract_t, internal::packet_traits&lt;RhsScalar&gt;::size,</div>
<div class="line"><a name="l00362"></a><span class="lineno">  362</span>&#160;        rhs_inner_dim_contiguous, rhs_inner_dim_reordered, <a class="codeRef" href="../group__enums.html#gga45fe06e29902b7a2773de05ba27b47a1a4e19dd09d5ff42295ba1d72d12a46686">Unaligned</a>&gt;</div>
<div class="line"><a name="l00363"></a><span class="lineno">  363</span>&#160;        RhsMapper;</div>
<div class="line"><a name="l00364"></a><span class="lineno">  364</span>&#160; </div>
<div class="line"><a name="l00365"></a><span class="lineno">  365</span>&#160;    <span class="keyword">typedef</span> internal::blas_data_mapper&lt;Scalar, Index, ColMajor&gt; OutputMapper;</div>
<div class="line"><a name="l00366"></a><span class="lineno">  366</span>&#160; </div>
<div class="line"><a name="l00367"></a><span class="lineno">  367</span>&#160;    <span class="keyword">typedef</span> internal::TensorContractionKernel&lt;</div>
<div class="line"><a name="l00368"></a><span class="lineno">  368</span>&#160;        Scalar, LhsScalar, RhsScalar, <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>, OutputMapper, LhsMapper, RhsMapper&gt;</div>
<div class="line"><a name="l00369"></a><span class="lineno">  369</span>&#160;        TensorContractionKernel;</div>
<div class="line"><a name="l00370"></a><span class="lineno">  370</span>&#160; </div>
<div class="line"><a name="l00371"></a><span class="lineno">  371</span>&#160;    <span class="keyword">typedef</span> <span class="keyword">typename</span> TensorContractionKernel::LhsBlock LhsBlock;</div>
<div class="line"><a name="l00372"></a><span class="lineno">  372</span>&#160;    <span class="keyword">typedef</span> <span class="keyword">typename</span> TensorContractionKernel::RhsBlock RhsBlock;</div>
<div class="line"><a name="l00373"></a><span class="lineno">  373</span>&#160;    <span class="keyword">typedef</span> <span class="keyword">typename</span> TensorContractionKernel::BlockMemHandle BlockMemHandle;</div>
<div class="line"><a name="l00374"></a><span class="lineno">  374</span>&#160; </div>
<div class="line"><a name="l00375"></a><span class="lineno">  375</span>&#160;    EvalParallelContext(<span class="keyword">const</span> Self* <span class="keyword">self</span>, <span class="keywordtype">int</span> num_threads, Scalar* buffer,</div>
<div class="line"><a name="l00376"></a><span class="lineno">  376</span>&#160;                        Index tm, Index tn, Index tk, Index bm, Index bn,</div>
<div class="line"><a name="l00377"></a><span class="lineno">  377</span>&#160;                        Index bk, Index nm, Index nn, Index nk, Index gm,</div>
<div class="line"><a name="l00378"></a><span class="lineno">  378</span>&#160;                        Index gn, Index nm0, Index nn0, <span class="keywordtype">bool</span> shard_by_col,</div>
<div class="line"><a name="l00379"></a><span class="lineno">  379</span>&#160;                        <span class="keywordtype">bool</span> parallel_pack,</div>
<div class="line"><a name="l00380"></a><span class="lineno">  380</span>&#160;                        <span class="keywordtype">bool</span> parallelize_by_sharding_dim_only,</div>
<div class="line"><a name="l00381"></a><span class="lineno">  381</span>&#160;                        DoneCallback done)</div>
<div class="line"><a name="l00382"></a><span class="lineno">  382</span>&#160;        : created_by_thread_id_(std::this_thread::get_id()),</div>
<div class="line"><a name="l00383"></a><span class="lineno">  383</span>&#160;          done_(this, std::move(done)),</div>
<div class="line"><a name="l00384"></a><span class="lineno">  384</span>&#160;          device_(self-&gt;m_device),</div>
<div class="line"><a name="l00385"></a><span class="lineno">  385</span>&#160;          lhs_(self-&gt;m_leftImpl, self-&gt;m_left_nocontract_strides,</div>
<div class="line"><a name="l00386"></a><span class="lineno">  386</span>&#160;               self-&gt;m_i_strides, self-&gt;m_left_contracting_strides,</div>
<div class="line"><a name="l00387"></a><span class="lineno">  387</span>&#160;               self-&gt;m_k_strides),</div>
<div class="line"><a name="l00388"></a><span class="lineno">  388</span>&#160;          rhs_(self-&gt;m_rightImpl, self-&gt;m_right_nocontract_strides,</div>
<div class="line"><a name="l00389"></a><span class="lineno">  389</span>&#160;               self-&gt;m_j_strides, self-&gt;m_right_contracting_strides,</div>
<div class="line"><a name="l00390"></a><span class="lineno">  390</span>&#160;               self-&gt;m_k_strides),</div>
<div class="line"><a name="l00391"></a><span class="lineno">  391</span>&#160;          buffer_(buffer),</div>
<div class="line"><a name="l00392"></a><span class="lineno">  392</span>&#160;          output_(buffer, tm),</div>
<div class="line"><a name="l00393"></a><span class="lineno">  393</span>&#160;          output_kernel_(self-&gt;m_output_kernel),</div>
<div class="line"><a name="l00394"></a><span class="lineno">  394</span>&#160;          tensor_contraction_params_(self-&gt;m_tensor_contraction_params),</div>
<div class="line"><a name="l00395"></a><span class="lineno">  395</span>&#160;          num_threads_(num_threads),</div>
<div class="line"><a name="l00396"></a><span class="lineno">  396</span>&#160;          shard_by_col_(shard_by_col),</div>
<div class="line"><a name="l00397"></a><span class="lineno">  397</span>&#160;          parallel_pack_(parallel_pack),</div>
<div class="line"><a name="l00398"></a><span class="lineno">  398</span>&#160;          parallelize_by_sharding_dim_only_(parallelize_by_sharding_dim_only),</div>
<div class="line"><a name="l00399"></a><span class="lineno">  399</span>&#160;          m_(tm),</div>
<div class="line"><a name="l00400"></a><span class="lineno">  400</span>&#160;          n_(tn),</div>
<div class="line"><a name="l00401"></a><span class="lineno">  401</span>&#160;          k_(tk),</div>
<div class="line"><a name="l00402"></a><span class="lineno">  402</span>&#160;          bm_(bm),</div>
<div class="line"><a name="l00403"></a><span class="lineno">  403</span>&#160;          bn_(bn),</div>
<div class="line"><a name="l00404"></a><span class="lineno">  404</span>&#160;          bk_(bk),</div>
<div class="line"><a name="l00405"></a><span class="lineno">  405</span>&#160;          nm_(nm),</div>
<div class="line"><a name="l00406"></a><span class="lineno">  406</span>&#160;          nn_(nn),</div>
<div class="line"><a name="l00407"></a><span class="lineno">  407</span>&#160;          nk_(nk),</div>
<div class="line"><a name="l00408"></a><span class="lineno">  408</span>&#160;          gm_(gm),</div>
<div class="line"><a name="l00409"></a><span class="lineno">  409</span>&#160;          gn_(gn),</div>
<div class="line"><a name="l00410"></a><span class="lineno">  410</span>&#160;          nm0_(nm0),</div>
<div class="line"><a name="l00411"></a><span class="lineno">  411</span>&#160;          nn0_(nn0),</div>
<div class="line"><a name="l00412"></a><span class="lineno">  412</span>&#160;          kernel_(m_, k_, n_, bm_, bk_, bn_),</div>
<div class="line"><a name="l00413"></a><span class="lineno">  413</span>&#160;          num_thread_local_allocations_(0),</div>
<div class="line"><a name="l00414"></a><span class="lineno">  414</span>&#160;          <span class="comment">// We reserve 2X more capacity for a thread local values, than the</span></div>
<div class="line"><a name="l00415"></a><span class="lineno">  415</span>&#160;          <span class="comment">// number of threads in the pool to efficiently handle task stealing</span></div>
<div class="line"><a name="l00416"></a><span class="lineno">  416</span>&#160;          <span class="comment">// by threads that are not managed by the pool.</span></div>
<div class="line"><a name="l00417"></a><span class="lineno">  417</span>&#160;          thread_local_capacity(2 * (parallelize_by_sharding_dim_only_</div>
<div class="line"><a name="l00418"></a><span class="lineno">  418</span>&#160;                                         ? device_.numThreadsInPool()</div>
<div class="line"><a name="l00419"></a><span class="lineno">  419</span>&#160;                                         : 0)),</div>
<div class="line"><a name="l00420"></a><span class="lineno">  420</span>&#160;          <span class="comment">// We will use only one of the Lhs/Rhs thread local storage depending</span></div>
<div class="line"><a name="l00421"></a><span class="lineno">  421</span>&#160;          <span class="comment">// on the shard_by_col value and we parallelize by sharding dim ONLY.</span></div>
<div class="line"><a name="l00422"></a><span class="lineno">  422</span>&#160;          lhs_thread_local_blocks_(shard_by_col_ ? 0 : thread_local_capacity,</div>
<div class="line"><a name="l00423"></a><span class="lineno">  423</span>&#160;                                   {*<span class="keyword">this</span>}, {*<span class="keyword">this</span>}),</div>
<div class="line"><a name="l00424"></a><span class="lineno">  424</span>&#160;          rhs_thread_local_blocks_(shard_by_col_ ? thread_local_capacity : 0,</div>
<div class="line"><a name="l00425"></a><span class="lineno">  425</span>&#160;                                   {*<span class="keyword">this</span>}, {*<span class="keyword">this</span>}) {</div>
<div class="line"><a name="l00426"></a><span class="lineno">  426</span>&#160;      <span class="comment">// These two options are mutually exclusive.</span></div>
<div class="line"><a name="l00427"></a><span class="lineno">  427</span>&#160;      eigen_assert(!(parallel_pack &amp;&amp; parallelize_by_sharding_dim_only));</div>
<div class="line"><a name="l00428"></a><span class="lineno">  428</span>&#160; </div>
<div class="line"><a name="l00429"></a><span class="lineno">  429</span>&#160;      <span class="keywordflow">for</span> (Index x = 0; x &lt; P; x++) {</div>
<div class="line"><a name="l00430"></a><span class="lineno">  430</span>&#160;        <span class="comment">// Normal number of notifications for k slice switch is</span></div>
<div class="line"><a name="l00431"></a><span class="lineno">  431</span>&#160;        <span class="comment">// nm_ + nn_ + nm_ * nn_. However, first P - 1 slices will receive only</span></div>
<div class="line"><a name="l00432"></a><span class="lineno">  432</span>&#160;        <span class="comment">// nm_ + nn_ notifications, because they will not receive notifications</span></div>
<div class="line"><a name="l00433"></a><span class="lineno">  433</span>&#160;        <span class="comment">// from preceding kernels.</span></div>
<div class="line"><a name="l00434"></a><span class="lineno">  434</span>&#160;        state_switch_[x] =</div>
<div class="line"><a name="l00435"></a><span class="lineno">  435</span>&#160;            x == 0</div>
<div class="line"><a name="l00436"></a><span class="lineno">  436</span>&#160;                ? 1</div>
<div class="line"><a name="l00437"></a><span class="lineno">  437</span>&#160;                : (parallel_pack_ ? nn_ + nm_ : (shard_by_col_ ? nn_ : nm_)) +</div>
<div class="line"><a name="l00438"></a><span class="lineno">  438</span>&#160;                      (x == P - 1 ? nm_ * nn_ : 0);</div>
<div class="line"><a name="l00439"></a><span class="lineno">  439</span>&#160;        state_packing_ready_[x] =</div>
<div class="line"><a name="l00440"></a><span class="lineno">  440</span>&#160;            parallel_pack_ ? 0 : (shard_by_col_ ? nm_ : nn_);</div>
<div class="line"><a name="l00441"></a><span class="lineno">  441</span>&#160;        state_kernel_[x] = <span class="keyword">new</span> std::atomic&lt;uint8_t&gt;*[nm_];</div>
<div class="line"><a name="l00442"></a><span class="lineno">  442</span>&#160;        <span class="keywordflow">for</span> (Index m = 0; m &lt; nm_; m++) {</div>
<div class="line"><a name="l00443"></a><span class="lineno">  443</span>&#160;          state_kernel_[x][m] = <span class="keyword">new</span> std::atomic&lt;uint8_t&gt;[nn_];</div>
<div class="line"><a name="l00444"></a><span class="lineno">  444</span>&#160;          <span class="comment">// Kernels generally receive 3 notifications (previous kernel + 2</span></div>
<div class="line"><a name="l00445"></a><span class="lineno">  445</span>&#160;          <span class="comment">// packing), but the first slice won&#39;t get notifications from previous</span></div>
<div class="line"><a name="l00446"></a><span class="lineno">  446</span>&#160;          <span class="comment">// kernels.</span></div>
<div class="line"><a name="l00447"></a><span class="lineno">  447</span>&#160;          <span class="keywordflow">for</span> (Index n = 0; n &lt; nn_; n++)</div>
<div class="line"><a name="l00448"></a><span class="lineno">  448</span>&#160;            state_kernel_[x][m][n].store(</div>
<div class="line"><a name="l00449"></a><span class="lineno">  449</span>&#160;                (x == 0 ? 0 : 1) + (parallel_pack_ ? 2 : 1),</div>
<div class="line"><a name="l00450"></a><span class="lineno">  450</span>&#160;                std::memory_order_relaxed);</div>
<div class="line"><a name="l00451"></a><span class="lineno">  451</span>&#160;        }</div>
<div class="line"><a name="l00452"></a><span class="lineno">  452</span>&#160;      }</div>
<div class="line"><a name="l00453"></a><span class="lineno">  453</span>&#160; </div>
<div class="line"><a name="l00454"></a><span class="lineno">  454</span>&#160;      <span class="comment">// Allocate memory for packed rhs/lhs matrices.</span></div>
<div class="line"><a name="l00455"></a><span class="lineno">  455</span>&#160;      packed_mem_ = kernel_.allocateSlices(            <span class="comment">//</span></div>
<div class="line"><a name="l00456"></a><span class="lineno">  456</span>&#160;          device_,                                     <span class="comment">//</span></div>
<div class="line"><a name="l00457"></a><span class="lineno">  457</span>&#160;          <span class="comment">/*num_lhs=*/</span>nm0_,                            <span class="comment">//</span></div>
<div class="line"><a name="l00458"></a><span class="lineno">  458</span>&#160;          <span class="comment">/*num_rhs=*/</span>nn0_,                            <span class="comment">//</span></div>
<div class="line"><a name="l00459"></a><span class="lineno">  459</span>&#160;          <span class="comment">/*num_slices=*/</span>std::min&lt;Index&gt;(nk_, P - 1),  <span class="comment">//</span></div>
<div class="line"><a name="l00460"></a><span class="lineno">  460</span>&#160;          packed_lhs_, packed_rhs_);</div>
<div class="line"><a name="l00461"></a><span class="lineno">  461</span>&#160; </div>
<div class="line"><a name="l00462"></a><span class="lineno">  462</span>&#160;      <span class="keywordflow">if</span> (parallelize_by_sharding_dim_only_) {</div>
<div class="line"><a name="l00463"></a><span class="lineno">  463</span>&#160;        <span class="keyword">const</span> <span class="keywordtype">int</span> num_worker_threads = device_.numThreadsInPool();</div>
<div class="line"><a name="l00464"></a><span class="lineno">  464</span>&#160; </div>
<div class="line"><a name="l00465"></a><span class="lineno">  465</span>&#160;        <span class="keywordflow">if</span> (shard_by_col) {</div>
<div class="line"><a name="l00466"></a><span class="lineno">  466</span>&#160;          can_use_thread_local_packed_ = <span class="keyword">new</span> std::atomic&lt;bool&gt;[nn_];</div>
<div class="line"><a name="l00467"></a><span class="lineno">  467</span>&#160;          <span class="keywordflow">for</span> (<span class="keywordtype">int</span> i = 0; i &lt; nn_; ++i)</div>
<div class="line"><a name="l00468"></a><span class="lineno">  468</span>&#160;            can_use_thread_local_packed_[i].store(<span class="keyword">true</span>,</div>
<div class="line"><a name="l00469"></a><span class="lineno">  469</span>&#160;                                                  std::memory_order_relaxed);</div>
<div class="line"><a name="l00470"></a><span class="lineno">  470</span>&#160; </div>
<div class="line"><a name="l00471"></a><span class="lineno">  471</span>&#160;          <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> num_blocks = num_worker_threads * gn_;</div>
<div class="line"><a name="l00472"></a><span class="lineno">  472</span>&#160;          thread_local_pre_alocated_mem_ = kernel_.allocateSlices(  <span class="comment">//</span></div>
<div class="line"><a name="l00473"></a><span class="lineno">  473</span>&#160;              device_,                                              <span class="comment">//</span></div>
<div class="line"><a name="l00474"></a><span class="lineno">  474</span>&#160;              <span class="comment">/*num_lhs=*/</span>0,                                        <span class="comment">//</span></div>
<div class="line"><a name="l00475"></a><span class="lineno">  475</span>&#160;              <span class="comment">/*num_rhs=*/</span>num_blocks,                               <span class="comment">//</span></div>
<div class="line"><a name="l00476"></a><span class="lineno">  476</span>&#160;              <span class="comment">/*num_slices=*/</span>1,                                     <span class="comment">//</span></div>
<div class="line"><a name="l00477"></a><span class="lineno">  477</span>&#160;              <span class="comment">/*lhs_blocks=*/</span><span class="keyword">nullptr</span>, &amp;rhs_thread_local_pre_allocated_);</div>
<div class="line"><a name="l00478"></a><span class="lineno">  478</span>&#160; </div>
<div class="line"><a name="l00479"></a><span class="lineno">  479</span>&#160;        } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00480"></a><span class="lineno">  480</span>&#160;          can_use_thread_local_packed_ = <span class="keyword">new</span> std::atomic&lt;bool&gt;[nm_];</div>
<div class="line"><a name="l00481"></a><span class="lineno">  481</span>&#160;          <span class="keywordflow">for</span> (<span class="keywordtype">int</span> i = 0; i &lt; nm_; ++i)</div>
<div class="line"><a name="l00482"></a><span class="lineno">  482</span>&#160;            can_use_thread_local_packed_[i].store(<span class="keyword">true</span>,</div>
<div class="line"><a name="l00483"></a><span class="lineno">  483</span>&#160;                                                  std::memory_order_relaxed);</div>
<div class="line"><a name="l00484"></a><span class="lineno">  484</span>&#160; </div>
<div class="line"><a name="l00485"></a><span class="lineno">  485</span>&#160;          <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> num_blocks = num_worker_threads * gm_;</div>
<div class="line"><a name="l00486"></a><span class="lineno">  486</span>&#160;          thread_local_pre_alocated_mem_ = kernel_.allocateSlices(  <span class="comment">//</span></div>
<div class="line"><a name="l00487"></a><span class="lineno">  487</span>&#160;              device_,                                              <span class="comment">//</span></div>
<div class="line"><a name="l00488"></a><span class="lineno">  488</span>&#160;              <span class="comment">/*num_lhs=*/</span>num_blocks,                               <span class="comment">//</span></div>
<div class="line"><a name="l00489"></a><span class="lineno">  489</span>&#160;              <span class="comment">/*num_rhs=*/</span>0,                                        <span class="comment">//</span></div>
<div class="line"><a name="l00490"></a><span class="lineno">  490</span>&#160;              <span class="comment">/*num_slices=*/</span>1, &amp;lhs_thread_local_pre_allocated_,   <span class="comment">//</span></div>
<div class="line"><a name="l00491"></a><span class="lineno">  491</span>&#160;              <span class="comment">/*rhs_blocks=*/</span><span class="keyword">nullptr</span>);</div>
<div class="line"><a name="l00492"></a><span class="lineno">  492</span>&#160;        }</div>
<div class="line"><a name="l00493"></a><span class="lineno">  493</span>&#160;      }</div>
<div class="line"><a name="l00494"></a><span class="lineno">  494</span>&#160;    }</div>
<div class="line"><a name="l00495"></a><span class="lineno">  495</span>&#160; </div>
<div class="line"><a name="l00496"></a><span class="lineno">  496</span>&#160;    ~EvalParallelContext() {</div>
<div class="line"><a name="l00497"></a><span class="lineno">  497</span>&#160;      <span class="keywordflow">for</span> (Index x = 0; x &lt; P; x++) {</div>
<div class="line"><a name="l00498"></a><span class="lineno">  498</span>&#160;        <span class="keywordflow">for</span> (Index m = 0; m &lt; nm_; m++) <span class="keyword">delete</span>[] state_kernel_[x][m];</div>
<div class="line"><a name="l00499"></a><span class="lineno">  499</span>&#160;        <span class="keyword">delete</span>[] state_kernel_[x];</div>
<div class="line"><a name="l00500"></a><span class="lineno">  500</span>&#160;      }</div>
<div class="line"><a name="l00501"></a><span class="lineno">  501</span>&#160;      kernel_.deallocate(device_, packed_mem_);</div>
<div class="line"><a name="l00502"></a><span class="lineno">  502</span>&#160;      <span class="keywordflow">if</span> (parallelize_by_sharding_dim_only_) {</div>
<div class="line"><a name="l00503"></a><span class="lineno">  503</span>&#160;        kernel_.deallocate(device_, thread_local_pre_alocated_mem_);</div>
<div class="line"><a name="l00504"></a><span class="lineno">  504</span>&#160;        <span class="keyword">delete</span>[] can_use_thread_local_packed_;</div>
<div class="line"><a name="l00505"></a><span class="lineno">  505</span>&#160;      }</div>
<div class="line"><a name="l00506"></a><span class="lineno">  506</span>&#160;    }</div>
<div class="line"><a name="l00507"></a><span class="lineno">  507</span>&#160; </div>
<div class="line"><a name="l00508"></a><span class="lineno">  508</span>&#160;    <span class="keywordtype">void</span> run() {</div>
<div class="line"><a name="l00509"></a><span class="lineno">  509</span>&#160;      <span class="comment">// Kick off packing of the first slice.</span></div>
<div class="line"><a name="l00510"></a><span class="lineno">  510</span>&#160;      signal_switch(0, 1);</div>
<div class="line"><a name="l00511"></a><span class="lineno">  511</span>&#160; </div>
<div class="line"><a name="l00512"></a><span class="lineno">  512</span>&#160;      <span class="comment">// Wait for overall completion.</span></div>
<div class="line"><a name="l00513"></a><span class="lineno">  513</span>&#160;      <span class="comment">//</span></div>
<div class="line"><a name="l00514"></a><span class="lineno">  514</span>&#160;      <span class="comment">// If parallel evaluation is executed in async mode, this is a no-op, and</span></div>
<div class="line"><a name="l00515"></a><span class="lineno">  515</span>&#160;      <span class="comment">// Wait() will return immediately. In synchronous mode it will block the</span></div>
<div class="line"><a name="l00516"></a><span class="lineno">  516</span>&#160;      <span class="comment">// caller thread until it will receive notification from last task.</span></div>
<div class="line"><a name="l00517"></a><span class="lineno">  517</span>&#160;      <span class="comment">//</span></div>
<div class="line"><a name="l00518"></a><span class="lineno">  518</span>&#160;      <span class="comment">// In async mode, last task when completed will call done callback from</span></div>
<div class="line"><a name="l00519"></a><span class="lineno">  519</span>&#160;      <span class="comment">// the same thread, and will delete this context.</span></div>
<div class="line"><a name="l00520"></a><span class="lineno">  520</span>&#160;      <span class="comment">//</span></div>
<div class="line"><a name="l00521"></a><span class="lineno">  521</span>&#160;      <span class="comment">// TODO(dvyukov): This wait can lead to deadlock if contraction is</span></div>
<div class="line"><a name="l00522"></a><span class="lineno">  522</span>&#160;      <span class="comment">// evaluated in synchronous mode. If nthreads contractions are</span></div>
<div class="line"><a name="l00523"></a><span class="lineno">  523</span>&#160;      <span class="comment">// concurrently submitted from worker threads, this wait will block all</span></div>
<div class="line"><a name="l00524"></a><span class="lineno">  524</span>&#160;      <span class="comment">// worker threads and the system will deadlock.</span></div>
<div class="line"><a name="l00525"></a><span class="lineno">  525</span>&#160;      done_.Wait();</div>
<div class="line"><a name="l00526"></a><span class="lineno">  526</span>&#160;    }</div>
<div class="line"><a name="l00527"></a><span class="lineno">  527</span>&#160; </div>
<div class="line"><a name="l00528"></a><span class="lineno">  528</span>&#160;   <span class="keyword">private</span>:</div>
<div class="line"><a name="l00529"></a><span class="lineno">  529</span>&#160;    std::thread::id created_by_thread_id_;</div>
<div class="line"><a name="l00530"></a><span class="lineno">  530</span>&#160; </div>
<div class="line"><a name="l00531"></a><span class="lineno">  531</span>&#160;    <span class="comment">// This notification is specialized on the type of DoneCallback and can be</span></div>
<div class="line"><a name="l00532"></a><span class="lineno">  532</span>&#160;    <span class="comment">// blocking or non-blocking.</span></div>
<div class="line"><a name="l00533"></a><span class="lineno">  533</span>&#160;    EvalParallelNotification&lt;DoneCallback, EvalParallelContext&gt; done_;</div>
<div class="line"><a name="l00534"></a><span class="lineno">  534</span>&#160; </div>
<div class="line"><a name="l00535"></a><span class="lineno">  535</span>&#160;    <span class="keyword">const</span> Device&amp; device_;</div>
<div class="line"><a name="l00536"></a><span class="lineno">  536</span>&#160;    LhsMapper lhs_;</div>
<div class="line"><a name="l00537"></a><span class="lineno">  537</span>&#160;    RhsMapper rhs_;</div>
<div class="line"><a name="l00538"></a><span class="lineno">  538</span>&#160;    Scalar* <span class="keyword">const</span> buffer_;</div>
<div class="line"><a name="l00539"></a><span class="lineno">  539</span>&#160;    OutputMapper output_;</div>
<div class="line"><a name="l00540"></a><span class="lineno">  540</span>&#160;    OutputKernelType output_kernel_;</div>
<div class="line"><a name="l00541"></a><span class="lineno">  541</span>&#160;    TensorContractionParams tensor_contraction_params_;</div>
<div class="line"><a name="l00542"></a><span class="lineno">  542</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">int</span> num_threads_;</div>
<div class="line"><a name="l00543"></a><span class="lineno">  543</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">bool</span> shard_by_col_;</div>
<div class="line"><a name="l00544"></a><span class="lineno">  544</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">bool</span> parallel_pack_;</div>
<div class="line"><a name="l00545"></a><span class="lineno">  545</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">bool</span> parallelize_by_sharding_dim_only_;</div>
<div class="line"><a name="l00546"></a><span class="lineno">  546</span>&#160;    <span class="comment">// Matrix sizes.</span></div>
<div class="line"><a name="l00547"></a><span class="lineno">  547</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> m_;</div>
<div class="line"><a name="l00548"></a><span class="lineno">  548</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> n_;</div>
<div class="line"><a name="l00549"></a><span class="lineno">  549</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> k_;</div>
<div class="line"><a name="l00550"></a><span class="lineno">  550</span>&#160;    <span class="comment">// Block sizes.</span></div>
<div class="line"><a name="l00551"></a><span class="lineno">  551</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> bm_;</div>
<div class="line"><a name="l00552"></a><span class="lineno">  552</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> bn_;</div>
<div class="line"><a name="l00553"></a><span class="lineno">  553</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> bk_;</div>
<div class="line"><a name="l00554"></a><span class="lineno">  554</span>&#160;    <span class="comment">// Number of tasks.</span></div>
<div class="line"><a name="l00555"></a><span class="lineno">  555</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nm_;</div>
<div class="line"><a name="l00556"></a><span class="lineno">  556</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nn_;</div>
<div class="line"><a name="l00557"></a><span class="lineno">  557</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nk_;</div>
<div class="line"><a name="l00558"></a><span class="lineno">  558</span>&#160;    <span class="comment">// Task grain sizes (number of kernels executed per task).</span></div>
<div class="line"><a name="l00559"></a><span class="lineno">  559</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> gm_;</div>
<div class="line"><a name="l00560"></a><span class="lineno">  560</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> gn_;</div>
<div class="line"><a name="l00561"></a><span class="lineno">  561</span>&#160;    <span class="comment">// Number of blocks (this is different from ni_/nn_ because of task size</span></div>
<div class="line"><a name="l00562"></a><span class="lineno">  562</span>&#160;    <span class="comment">// coarsening).</span></div>
<div class="line"><a name="l00563"></a><span class="lineno">  563</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nm0_;</div>
<div class="line"><a name="l00564"></a><span class="lineno">  564</span>&#160;    <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nn0_;</div>
<div class="line"><a name="l00565"></a><span class="lineno">  565</span>&#160;    <span class="comment">// Tensor contraction kernel.</span></div>
<div class="line"><a name="l00566"></a><span class="lineno">  566</span>&#160;    TensorContractionKernel kernel_;</div>
<div class="line"><a name="l00567"></a><span class="lineno">  567</span>&#160; </div>
<div class="line"><a name="l00568"></a><span class="lineno">  568</span>&#160;    <span class="comment">// Parallelization strategy.</span></div>
<div class="line"><a name="l00569"></a><span class="lineno">  569</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00570"></a><span class="lineno">  570</span>&#160;    <span class="comment">// Blocks related to the same k block can run in parallel because they write</span></div>
<div class="line"><a name="l00571"></a><span class="lineno">  571</span>&#160;    <span class="comment">// to different output blocks. So we parallelize within k slices, this</span></div>
<div class="line"><a name="l00572"></a><span class="lineno">  572</span>&#160;    <span class="comment">// gives us parallelism level of m x n. Before we can start any kernels</span></div>
<div class="line"><a name="l00573"></a><span class="lineno">  573</span>&#160;    <span class="comment">// related to k-th slice, we need to issue m lhs packing tasks and n rhs</span></div>
<div class="line"><a name="l00574"></a><span class="lineno">  574</span>&#160;    <span class="comment">// packing tasks.</span></div>
<div class="line"><a name="l00575"></a><span class="lineno">  575</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00576"></a><span class="lineno">  576</span>&#160;    <span class="comment">// However, there is a bottleneck when we are finishing kernels for k-th</span></div>
<div class="line"><a name="l00577"></a><span class="lineno">  577</span>&#160;    <span class="comment">// slice (at the very end there is only 1 runnable kernel). To mitigate this</span></div>
<div class="line"><a name="l00578"></a><span class="lineno">  578</span>&#160;    <span class="comment">// bottleneck we allow kernels from k-th and k+1-th slices to run in</span></div>
<div class="line"><a name="l00579"></a><span class="lineno">  579</span>&#160;    <span class="comment">// parallel. Note that (m, n, k) and (m, n, k+1) kernels write to the same</span></div>
<div class="line"><a name="l00580"></a><span class="lineno">  580</span>&#160;    <span class="comment">// output block, so they must not run in parallel.</span></div>
<div class="line"><a name="l00581"></a><span class="lineno">  581</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00582"></a><span class="lineno">  582</span>&#160;    <span class="comment">// This gives us the following dependency graph.</span></div>
<div class="line"><a name="l00583"></a><span class="lineno">  583</span>&#160;    <span class="comment">// On each k slice we have m x n kernel tasks, m lhs paking tasks and n rhs</span></div>
<div class="line"><a name="l00584"></a><span class="lineno">  584</span>&#160;    <span class="comment">// packing tasks.</span></div>
<div class="line"><a name="l00585"></a><span class="lineno">  585</span>&#160;    <span class="comment">// Kernel (m, n, k) can start when:</span></div>
<div class="line"><a name="l00586"></a><span class="lineno">  586</span>&#160;    <span class="comment">//  - kernel (m, n, k-1) has finished</span></div>
<div class="line"><a name="l00587"></a><span class="lineno">  587</span>&#160;    <span class="comment">//  - lhs packing (m, k) has finished</span></div>
<div class="line"><a name="l00588"></a><span class="lineno">  588</span>&#160;    <span class="comment">//  - rhs packing (n, k) has finished</span></div>
<div class="line"><a name="l00589"></a><span class="lineno">  589</span>&#160;    <span class="comment">// Lhs/rhs packing can start when:</span></div>
<div class="line"><a name="l00590"></a><span class="lineno">  590</span>&#160;    <span class="comment">//  - all k-1 packing has finished (artificially imposed to limit amount of</span></div>
<div class="line"><a name="l00591"></a><span class="lineno">  591</span>&#160;    <span class="comment">//  parallel packing)</span></div>
<div class="line"><a name="l00592"></a><span class="lineno">  592</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00593"></a><span class="lineno">  593</span>&#160;    <span class="comment">// On top of that we limit runnable tasks to two consecutive k slices.</span></div>
<div class="line"><a name="l00594"></a><span class="lineno">  594</span>&#160;    <span class="comment">// This is done to limit amount of memory we need for packed lhs/rhs</span></div>
<div class="line"><a name="l00595"></a><span class="lineno">  595</span>&#160;    <span class="comment">// (for each k slice we need m*bk + n*bk memory in packed_lhs_/packed_rhs_).</span></div>
<div class="line"><a name="l00596"></a><span class="lineno">  596</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00597"></a><span class="lineno">  597</span>&#160;    <span class="comment">// state_switch_ tracks when we are ready to switch to the next k slice.</span></div>
<div class="line"><a name="l00598"></a><span class="lineno">  598</span>&#160;    <span class="comment">// state_kernel_[m][n] tracks when we are ready to kick off kernel (m, n).</span></div>
<div class="line"><a name="l00599"></a><span class="lineno">  599</span>&#160;    <span class="comment">// These variable are rolling over 3 consecutive k slices: first two we are</span></div>
<div class="line"><a name="l00600"></a><span class="lineno">  600</span>&#160;    <span class="comment">// actively executing + one to track completion of kernels in the second</span></div>
<div class="line"><a name="l00601"></a><span class="lineno">  601</span>&#160;    <span class="comment">// slice.</span></div>
<div class="line"><a name="l00602"></a><span class="lineno">  602</span>&#160;    <span class="keyword">static</span> constexpr <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> P = 3;</div>
<div class="line"><a name="l00603"></a><span class="lineno">  603</span>&#160; </div>
<div class="line"><a name="l00604"></a><span class="lineno">  604</span>&#160;    <span class="comment">// Handle to the allocated temporary storage for Lhs/Rhs blocks.</span></div>
<div class="line"><a name="l00605"></a><span class="lineno">  605</span>&#160;    BlockMemHandle packed_mem_;</div>
<div class="line"><a name="l00606"></a><span class="lineno">  606</span>&#160;    std::vector&lt;LhsBlock&gt; packed_lhs_[P - 1];</div>
<div class="line"><a name="l00607"></a><span class="lineno">  607</span>&#160;    std::vector&lt;RhsBlock&gt; packed_rhs_[P - 1];</div>
<div class="line"><a name="l00608"></a><span class="lineno">  608</span>&#160; </div>
<div class="line"><a name="l00609"></a><span class="lineno">  609</span>&#160;    <span class="comment">// If we choose to parallelize only by the sharding dimension, each thread</span></div>
<div class="line"><a name="l00610"></a><span class="lineno">  610</span>&#160;    <span class="comment">// will have it&#39;s own &quot;thead local&quot; (not a c++ thread local storage) memory</span></div>
<div class="line"><a name="l00611"></a><span class="lineno">  611</span>&#160;    <span class="comment">// for packed_lhs or packed_rhs (shard_by_col = false of true). This memory</span></div>
<div class="line"><a name="l00612"></a><span class="lineno">  612</span>&#160;    <span class="comment">// can&#39;t be passed to a kernel that might execute on a different thread.</span></div>
<div class="line"><a name="l00613"></a><span class="lineno">  613</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00614"></a><span class="lineno">  614</span>&#160;    <span class="comment">// In practice when we are ready to pack memory for the sharding dimension</span></div>
<div class="line"><a name="l00615"></a><span class="lineno">  615</span>&#160;    <span class="comment">// (rhs if shard_by_col==true) of the K-th slice, all kernels for K-1 slice</span></div>
<div class="line"><a name="l00616"></a><span class="lineno">  616</span>&#160;    <span class="comment">// already computed (99% of the time), and we can pack data into the thread</span></div>
<div class="line"><a name="l00617"></a><span class="lineno">  617</span>&#160;    <span class="comment">// local storage, and guarantee that all the kernels will be executed</span></div>
<div class="line"><a name="l00618"></a><span class="lineno">  618</span>&#160;    <span class="comment">// immediately in the same thread. This significantly increases L1 cache hit</span></div>
<div class="line"><a name="l00619"></a><span class="lineno">  619</span>&#160;    <span class="comment">// ratio and reduces pressure on the memory bus.</span></div>
<div class="line"><a name="l00620"></a><span class="lineno">  620</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00621"></a><span class="lineno">  621</span>&#160;    <span class="comment">// It&#39;s still possible that kernel for the K-th slice will be ready before</span></div>
<div class="line"><a name="l00622"></a><span class="lineno">  622</span>&#160;    <span class="comment">// completion of the K-1 kernel, so we have to allocate &quot;global&quot; packed_lhs_</span></div>
<div class="line"><a name="l00623"></a><span class="lineno">  623</span>&#160;    <span class="comment">// and packed_rhs_ to allow kernels to be executed later on a thread</span></div>
<div class="line"><a name="l00624"></a><span class="lineno">  624</span>&#160;    <span class="comment">// different from the thread that was used for packing.</span></div>
<div class="line"><a name="l00625"></a><span class="lineno">  625</span>&#160; </div>
<div class="line"><a name="l00626"></a><span class="lineno">  626</span>&#160;    <span class="comment">// Handle for pre-allocated thread local memory buffers.</span></div>
<div class="line"><a name="l00627"></a><span class="lineno">  627</span>&#160;    BlockMemHandle thread_local_pre_alocated_mem_;</div>
<div class="line"><a name="l00628"></a><span class="lineno">  628</span>&#160; </div>
<div class="line"><a name="l00629"></a><span class="lineno">  629</span>&#160;    <span class="comment">// Only one of these will be initialized depending on shard_by_col value</span></div>
<div class="line"><a name="l00630"></a><span class="lineno">  630</span>&#160;    <span class="comment">// (the size will be `num_worker_threads * num_grains_in_the_sharding_dim`).</span></div>
<div class="line"><a name="l00631"></a><span class="lineno">  631</span>&#160;    std::vector&lt;LhsBlock&gt; lhs_thread_local_pre_allocated_;</div>
<div class="line"><a name="l00632"></a><span class="lineno">  632</span>&#160;    std::vector&lt;RhsBlock&gt; rhs_thread_local_pre_allocated_;</div>
<div class="line"><a name="l00633"></a><span class="lineno">  633</span>&#160; </div>
<div class="line"><a name="l00634"></a><span class="lineno">  634</span>&#160;    <span class="comment">// How many thread local blocks were already allocated.</span></div>
<div class="line"><a name="l00635"></a><span class="lineno">  635</span>&#160;    std::atomic&lt;int&gt; num_thread_local_allocations_;</div>
<div class="line"><a name="l00636"></a><span class="lineno">  636</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">int</span> thread_local_capacity;</div>
<div class="line"><a name="l00637"></a><span class="lineno">  637</span>&#160; </div>
<div class="line"><a name="l00638"></a><span class="lineno">  638</span>&#160;    <span class="comment">// We will use pre-allocated Lhs/Rhs blocks defined above, if the number of</span></div>
<div class="line"><a name="l00639"></a><span class="lineno">  639</span>&#160;    <span class="comment">// unique threads in a system is below or equal to the number of threads in</span></div>
<div class="line"><a name="l00640"></a><span class="lineno">  640</span>&#160;    <span class="comment">// a thread pool. We will fallback on dynamic memory allocation after that.</span></div>
<div class="line"><a name="l00641"></a><span class="lineno">  641</span>&#160; </div>
<div class="line"><a name="l00642"></a><span class="lineno">  642</span>&#160;    <span class="comment">// ThreadLocalBlocks is a container for Lhs or Rhs thread local buffers. Its</span></div>
<div class="line"><a name="l00643"></a><span class="lineno">  643</span>&#160;    <span class="comment">// size is equal to the grain size in Lhs/Rhs sharding dimension.</span></div>
<div class="line"><a name="l00644"></a><span class="lineno">  644</span>&#160;    <span class="keyword">template</span> &lt;<span class="keyword">typename</span> BlockType&gt;</div>
<div class="line"><a name="l00645"></a><span class="lineno">  645</span>&#160;    <span class="keyword">class </span>ThreadLocalBlocks {</div>
<div class="line"><a name="l00646"></a><span class="lineno">  646</span>&#160;     <span class="keyword">public</span>:</div>
<div class="line"><a name="l00647"></a><span class="lineno">  647</span>&#160;      ThreadLocalBlocks() = <span class="keywordflow">default</span>;</div>
<div class="line"><a name="l00648"></a><span class="lineno">  648</span>&#160; </div>
<div class="line"><a name="l00649"></a><span class="lineno">  649</span>&#160;      ThreadLocalBlocks(BlockType* base, <span class="keywordtype">size_t</span> grain_size)</div>
<div class="line"><a name="l00650"></a><span class="lineno">  650</span>&#160;          : is_pre_allocated_(true),</div>
<div class="line"><a name="l00651"></a><span class="lineno">  651</span>&#160;            thread_local_pre_allocated_base_(base),</div>
<div class="line"><a name="l00652"></a><span class="lineno">  652</span>&#160;            grain_size_(grain_size) {}</div>
<div class="line"><a name="l00653"></a><span class="lineno">  653</span>&#160; </div>
<div class="line"><a name="l00654"></a><span class="lineno">  654</span>&#160;      ThreadLocalBlocks(BlockMemHandle mem_handle,</div>
<div class="line"><a name="l00655"></a><span class="lineno">  655</span>&#160;                        std::vector&lt;BlockType&gt; blocks)</div>
<div class="line"><a name="l00656"></a><span class="lineno">  656</span>&#160;          : is_pre_allocated_(false),</div>
<div class="line"><a name="l00657"></a><span class="lineno">  657</span>&#160;            mem_handle_(std::move(mem_handle)),</div>
<div class="line"><a name="l00658"></a><span class="lineno">  658</span>&#160;            blocks_(std::move(blocks)) {}</div>
<div class="line"><a name="l00659"></a><span class="lineno">  659</span>&#160; </div>
<div class="line"><a name="l00660"></a><span class="lineno">  660</span>&#160;      BlockType&amp; block(<span class="keywordtype">int</span> grain_index) {</div>
<div class="line"><a name="l00661"></a><span class="lineno">  661</span>&#160;        eigen_assert(grain_index &gt;= 0);</div>
<div class="line"><a name="l00662"></a><span class="lineno">  662</span>&#160;        eigen_assert(<span class="keyword">static_cast&lt;</span><span class="keywordtype">size_t</span><span class="keyword">&gt;</span>(grain_index) &lt; size());</div>
<div class="line"><a name="l00663"></a><span class="lineno">  663</span>&#160;        <span class="keywordflow">return</span> is_pre_allocated_ ? thread_local_pre_allocated_base_[grain_index]</div>
<div class="line"><a name="l00664"></a><span class="lineno">  664</span>&#160;                                 : blocks_[grain_index];</div>
<div class="line"><a name="l00665"></a><span class="lineno">  665</span>&#160;      }</div>
<div class="line"><a name="l00666"></a><span class="lineno">  666</span>&#160; </div>
<div class="line"><a name="l00667"></a><span class="lineno">  667</span>&#160;      <span class="keywordtype">void</span> Release(EvalParallelContext&amp; ctx)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l00668"></a><span class="lineno">  668</span>&#160;        <span class="keywordflow">if</span> (!is_pre_allocated_) {</div>
<div class="line"><a name="l00669"></a><span class="lineno">  669</span>&#160;          ctx.kernel_.deallocate(ctx.device_, mem_handle_);</div>
<div class="line"><a name="l00670"></a><span class="lineno">  670</span>&#160;        }</div>
<div class="line"><a name="l00671"></a><span class="lineno">  671</span>&#160;      }</div>
<div class="line"><a name="l00672"></a><span class="lineno">  672</span>&#160; </div>
<div class="line"><a name="l00673"></a><span class="lineno">  673</span>&#160;      <span class="keywordtype">size_t</span> size()<span class="keyword"> const </span>{</div>
<div class="line"><a name="l00674"></a><span class="lineno">  674</span>&#160;        <span class="keywordflow">return</span> is_pre_allocated_ ? grain_size_ : blocks_.size();</div>
<div class="line"><a name="l00675"></a><span class="lineno">  675</span>&#160;      }</div>
<div class="line"><a name="l00676"></a><span class="lineno">  676</span>&#160; </div>
<div class="line"><a name="l00677"></a><span class="lineno">  677</span>&#160;     <span class="keyword">private</span>:</div>
<div class="line"><a name="l00678"></a><span class="lineno">  678</span>&#160;      <span class="keywordtype">bool</span> is_pre_allocated_;</div>
<div class="line"><a name="l00679"></a><span class="lineno">  679</span>&#160; </div>
<div class="line"><a name="l00680"></a><span class="lineno">  680</span>&#160;      <span class="comment">// Reuse pre-allocated thread local buffers.</span></div>
<div class="line"><a name="l00681"></a><span class="lineno">  681</span>&#160;      BlockType* thread_local_pre_allocated_base_ = <span class="keyword">nullptr</span>;</div>
<div class="line"><a name="l00682"></a><span class="lineno">  682</span>&#160;      <span class="keywordtype">size_t</span> grain_size_ = 0;</div>
<div class="line"><a name="l00683"></a><span class="lineno">  683</span>&#160; </div>
<div class="line"><a name="l00684"></a><span class="lineno">  684</span>&#160;      <span class="comment">// These will be initialized only if `is_pre_allocated == false`.</span></div>
<div class="line"><a name="l00685"></a><span class="lineno">  685</span>&#160;      BlockMemHandle mem_handle_{};</div>
<div class="line"><a name="l00686"></a><span class="lineno">  686</span>&#160;      std::vector&lt;BlockType&gt; blocks_;</div>
<div class="line"><a name="l00687"></a><span class="lineno">  687</span>&#160;    };</div>
<div class="line"><a name="l00688"></a><span class="lineno">  688</span>&#160; </div>
<div class="line"><a name="l00689"></a><span class="lineno">  689</span>&#160;    <span class="comment">// ThreadLocalBlocksInitialize callable does custom thread local blocks</span></div>
<div class="line"><a name="l00690"></a><span class="lineno">  690</span>&#160;    <span class="comment">// initialization, and will reuse pre-allocated buffers if possible, or will</span></div>
<div class="line"><a name="l00691"></a><span class="lineno">  691</span>&#160;    <span class="comment">// dynamically allocate new memory.</span></div>
<div class="line"><a name="l00692"></a><span class="lineno">  692</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00693"></a><span class="lineno">  693</span>&#160;    <span class="comment">// Lhs/Rhs blocks might be of the same type, so we have to pass explicitly</span></div>
<div class="line"><a name="l00694"></a><span class="lineno">  694</span>&#160;    <span class="comment">// for what side do we plan to do block allocation.</span></div>
<div class="line"><a name="l00695"></a><span class="lineno">  695</span>&#160;    <span class="keyword">template</span> &lt;<span class="keyword">typename</span> BlockType, <span class="keywordtype">bool</span> is_rhs&gt;</div>
<div class="line"><a name="l00696"></a><span class="lineno">  696</span>&#160;    <span class="keyword">class </span>ThreadLocalBlocksInitialize {</div>
<div class="line"><a name="l00697"></a><span class="lineno">  697</span>&#160;      <span class="keyword">static</span> constexpr <span class="keywordtype">bool</span> kIsLhs =</div>
<div class="line"><a name="l00698"></a><span class="lineno">  698</span>&#160;          !is_rhs &amp;&amp; std::is_same&lt;BlockType, LhsBlock&gt;::value;</div>
<div class="line"><a name="l00699"></a><span class="lineno">  699</span>&#160;      <span class="keyword">static</span> <span class="keyword">const</span> <span class="keywordtype">bool</span> kIsRhs =</div>
<div class="line"><a name="l00700"></a><span class="lineno">  700</span>&#160;          is_rhs &amp;&amp; std::is_same&lt;BlockType, RhsBlock&gt;::value;</div>
<div class="line"><a name="l00701"></a><span class="lineno">  701</span>&#160;      static_assert(kIsLhs || kIsRhs, <span class="stringliteral">&quot;Unknown block type&quot;</span>);</div>
<div class="line"><a name="l00702"></a><span class="lineno">  702</span>&#160; </div>
<div class="line"><a name="l00703"></a><span class="lineno">  703</span>&#160;      <span class="keyword">using</span> Blocks = ThreadLocalBlocks&lt;BlockType&gt;;</div>
<div class="line"><a name="l00704"></a><span class="lineno">  704</span>&#160; </div>
<div class="line"><a name="l00705"></a><span class="lineno">  705</span>&#160;     <span class="keyword">public</span>:</div>
<div class="line"><a name="l00706"></a><span class="lineno">  706</span>&#160;      ThreadLocalBlocksInitialize(EvalParallelContext&amp; ctx)</div>
<div class="line"><a name="l00707"></a><span class="lineno">  707</span>&#160;          : ctx_(ctx),</div>
<div class="line"><a name="l00708"></a><span class="lineno">  708</span>&#160;            num_worker_threads_(ctx_.device_.numThreadsInPool()) {}</div>
<div class="line"><a name="l00709"></a><span class="lineno">  709</span>&#160; </div>
<div class="line"><a name="l00710"></a><span class="lineno">  710</span>&#160;      <span class="keywordtype">void</span> operator()(Blocks&amp; blocks) {</div>
<div class="line"><a name="l00711"></a><span class="lineno">  711</span>&#160;        <span class="keyword">const</span> <span class="keywordtype">int</span> n = ctx_.num_thread_local_allocations_.fetch_add(</div>
<div class="line"><a name="l00712"></a><span class="lineno">  712</span>&#160;            1, std::memory_order_relaxed);</div>
<div class="line"><a name="l00713"></a><span class="lineno">  713</span>&#160; </div>
<div class="line"><a name="l00714"></a><span class="lineno">  714</span>&#160;        <span class="keywordflow">if</span> (n &gt;= num_worker_threads_) {</div>
<div class="line"><a name="l00715"></a><span class="lineno">  715</span>&#160;          ThreadLocalBlocksAllocator&lt;is_rhs&gt;::allocate(ctx_, blocks);</div>
<div class="line"><a name="l00716"></a><span class="lineno">  716</span>&#160;        } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00717"></a><span class="lineno">  717</span>&#160;          ThreadLocalBlocksAllocator&lt;is_rhs&gt;::reuse(ctx_, n, blocks);</div>
<div class="line"><a name="l00718"></a><span class="lineno">  718</span>&#160;        }</div>
<div class="line"><a name="l00719"></a><span class="lineno">  719</span>&#160;      }</div>
<div class="line"><a name="l00720"></a><span class="lineno">  720</span>&#160; </div>
<div class="line"><a name="l00721"></a><span class="lineno">  721</span>&#160;     <span class="keyword">private</span>:</div>
<div class="line"><a name="l00722"></a><span class="lineno">  722</span>&#160;      <span class="comment">// NOTE(ezhulenev): Without &#39;if constexpr&#39; we have to put calls to</span></div>
<div class="line"><a name="l00723"></a><span class="lineno">  723</span>&#160;      <span class="comment">// TensorContractionKernel::allocateSlices into template specializations.</span></div>
<div class="line"><a name="l00724"></a><span class="lineno">  724</span>&#160;      <span class="comment">// Also explicit specializations are not allowed at class scope in C++03,</span></div>
<div class="line"><a name="l00725"></a><span class="lineno">  725</span>&#160;      <span class="comment">// EvalCtx type parameter is just a workaround for that limitation.</span></div>
<div class="line"><a name="l00726"></a><span class="lineno">  726</span>&#160;      <span class="keyword">template</span> &lt;<span class="keywordtype">bool</span> pack_rhs, <span class="keyword">typename</span> EvalCtx = EvalParallelContext&gt;</div>
<div class="line"><a name="l00727"></a><span class="lineno">  727</span>&#160;      <span class="keyword">struct </span>ThreadLocalBlocksAllocator;</div>
<div class="line"><a name="l00728"></a><span class="lineno">  728</span>&#160; </div>
<div class="line"><a name="l00729"></a><span class="lineno">  729</span>&#160;      <span class="keyword">template</span> &lt;<span class="keyword">typename</span> EvalCtx&gt;</div>
<div class="line"><a name="l00730"></a><span class="lineno">  730</span>&#160;      <span class="keyword">struct </span>ThreadLocalBlocksAllocator&lt;<span class="comment">/*pack_rhs=*/</span>true, EvalCtx&gt; {</div>
<div class="line"><a name="l00731"></a><span class="lineno">  731</span>&#160;        <span class="keyword">static</span> <span class="keywordtype">void</span> allocate(EvalCtx&amp; ctx, Blocks&amp; blocks) {</div>
<div class="line"><a name="l00732"></a><span class="lineno">  732</span>&#160;          std::vector&lt;RhsBlock&gt; rhs_blocks;</div>
<div class="line"><a name="l00733"></a><span class="lineno">  733</span>&#160;          BlockMemHandle mem_handle = ctx.kernel_.allocateSlices(</div>
<div class="line"><a name="l00734"></a><span class="lineno">  734</span>&#160;              ctx.device_,</div>
<div class="line"><a name="l00735"></a><span class="lineno">  735</span>&#160;              <span class="comment">/*num_lhs=*/</span>0,</div>
<div class="line"><a name="l00736"></a><span class="lineno">  736</span>&#160;              <span class="comment">/*num_rhs=*/</span>ctx.gn_,</div>
<div class="line"><a name="l00737"></a><span class="lineno">  737</span>&#160;              <span class="comment">/*num_slices=*/</span>1,</div>
<div class="line"><a name="l00738"></a><span class="lineno">  738</span>&#160;              <span class="comment">/*lhs_blocks=*/</span><span class="keyword">nullptr</span>, <span class="comment">/*rhs_blocks=*/</span>&amp;rhs_blocks);</div>
<div class="line"><a name="l00739"></a><span class="lineno">  739</span>&#160; </div>
<div class="line"><a name="l00740"></a><span class="lineno">  740</span>&#160;          blocks = ThreadLocalBlocks&lt;RhsBlock&gt;(std::move(mem_handle),</div>
<div class="line"><a name="l00741"></a><span class="lineno">  741</span>&#160;                                               std::move(rhs_blocks));</div>
<div class="line"><a name="l00742"></a><span class="lineno">  742</span>&#160;        }</div>
<div class="line"><a name="l00743"></a><span class="lineno">  743</span>&#160; </div>
<div class="line"><a name="l00744"></a><span class="lineno">  744</span>&#160;        <span class="keyword">static</span> <span class="keywordtype">void</span> reuse(EvalCtx&amp; ctx, <span class="keywordtype">int</span> index, Blocks&amp; blocks) {</div>
<div class="line"><a name="l00745"></a><span class="lineno">  745</span>&#160;          RhsBlock* ptr = &amp;ctx.rhs_thread_local_pre_allocated_[ctx.gn_ * index];</div>
<div class="line"><a name="l00746"></a><span class="lineno">  746</span>&#160;          blocks = ThreadLocalBlocks&lt;RhsBlock&gt;(ptr, ctx.gn_);</div>
<div class="line"><a name="l00747"></a><span class="lineno">  747</span>&#160;        }</div>
<div class="line"><a name="l00748"></a><span class="lineno">  748</span>&#160;      };</div>
<div class="line"><a name="l00749"></a><span class="lineno">  749</span>&#160; </div>
<div class="line"><a name="l00750"></a><span class="lineno">  750</span>&#160;      <span class="keyword">template</span> &lt;<span class="keyword">typename</span> EvalCtx&gt;</div>
<div class="line"><a name="l00751"></a><span class="lineno">  751</span>&#160;      <span class="keyword">struct </span>ThreadLocalBlocksAllocator&lt;<span class="comment">/*pack_rhs=*/</span>false, EvalCtx&gt; {</div>
<div class="line"><a name="l00752"></a><span class="lineno">  752</span>&#160;        <span class="keyword">static</span> <span class="keywordtype">void</span> allocate(EvalCtx&amp; ctx, Blocks&amp; blocks) {</div>
<div class="line"><a name="l00753"></a><span class="lineno">  753</span>&#160;          std::vector&lt;LhsBlock&gt; lhs_blocks;</div>
<div class="line"><a name="l00754"></a><span class="lineno">  754</span>&#160;          BlockMemHandle mem_handle = ctx.kernel_.allocateSlices(</div>
<div class="line"><a name="l00755"></a><span class="lineno">  755</span>&#160;              ctx.device_,</div>
<div class="line"><a name="l00756"></a><span class="lineno">  756</span>&#160;              <span class="comment">/*num_lhs=*/</span>ctx.gm_,</div>
<div class="line"><a name="l00757"></a><span class="lineno">  757</span>&#160;              <span class="comment">/*num_rhs=*/</span>0,</div>
<div class="line"><a name="l00758"></a><span class="lineno">  758</span>&#160;              <span class="comment">/*num_slices=*/</span>1,</div>
<div class="line"><a name="l00759"></a><span class="lineno">  759</span>&#160;              <span class="comment">/*lhs_blocks=*/</span>&amp;lhs_blocks, <span class="comment">/*rhs_blocks=*/</span><span class="keyword">nullptr</span>);</div>
<div class="line"><a name="l00760"></a><span class="lineno">  760</span>&#160; </div>
<div class="line"><a name="l00761"></a><span class="lineno">  761</span>&#160;          blocks = ThreadLocalBlocks&lt;LhsBlock&gt;(std::move(mem_handle),</div>
<div class="line"><a name="l00762"></a><span class="lineno">  762</span>&#160;                                               std::move(lhs_blocks));</div>
<div class="line"><a name="l00763"></a><span class="lineno">  763</span>&#160;        }</div>
<div class="line"><a name="l00764"></a><span class="lineno">  764</span>&#160; </div>
<div class="line"><a name="l00765"></a><span class="lineno">  765</span>&#160;        <span class="keyword">static</span> <span class="keywordtype">void</span> reuse(EvalCtx&amp; ctx, <span class="keywordtype">int</span> index, Blocks&amp; blocks) {</div>
<div class="line"><a name="l00766"></a><span class="lineno">  766</span>&#160;          LhsBlock* ptr = &amp;ctx.lhs_thread_local_pre_allocated_[ctx.gm_ * index];</div>
<div class="line"><a name="l00767"></a><span class="lineno">  767</span>&#160;          blocks = ThreadLocalBlocks&lt;LhsBlock&gt;(ptr, ctx.gm_);</div>
<div class="line"><a name="l00768"></a><span class="lineno">  768</span>&#160;        }</div>
<div class="line"><a name="l00769"></a><span class="lineno">  769</span>&#160;      };</div>
<div class="line"><a name="l00770"></a><span class="lineno">  770</span>&#160; </div>
<div class="line"><a name="l00771"></a><span class="lineno">  771</span>&#160;      EvalParallelContext&amp; ctx_;</div>
<div class="line"><a name="l00772"></a><span class="lineno">  772</span>&#160;      <span class="keyword">const</span> <span class="keywordtype">int</span> num_worker_threads_;</div>
<div class="line"><a name="l00773"></a><span class="lineno">  773</span>&#160;    };</div>
<div class="line"><a name="l00774"></a><span class="lineno">  774</span>&#160; </div>
<div class="line"><a name="l00775"></a><span class="lineno">  775</span>&#160;    <span class="keyword">template</span> &lt;<span class="keyword">typename</span> BlockType&gt;</div>
<div class="line"><a name="l00776"></a><span class="lineno">  776</span>&#160;    <span class="keyword">class </span>ThreadLocalBlocksRelease {</div>
<div class="line"><a name="l00777"></a><span class="lineno">  777</span>&#160;     <span class="keyword">public</span>:</div>
<div class="line"><a name="l00778"></a><span class="lineno">  778</span>&#160;      <span class="keyword">using</span> Blocks = ThreadLocalBlocks&lt;BlockType&gt;;</div>
<div class="line"><a name="l00779"></a><span class="lineno">  779</span>&#160;      ThreadLocalBlocksRelease(EvalParallelContext&amp; ctx) : ctx_(ctx) {}</div>
<div class="line"><a name="l00780"></a><span class="lineno">  780</span>&#160;      <span class="keywordtype">void</span> operator()(Blocks&amp; blocks) { blocks.Release(ctx_); }</div>
<div class="line"><a name="l00781"></a><span class="lineno">  781</span>&#160; </div>
<div class="line"><a name="l00782"></a><span class="lineno">  782</span>&#160;     <span class="keyword">private</span>:</div>
<div class="line"><a name="l00783"></a><span class="lineno">  783</span>&#160;      EvalParallelContext&amp; ctx_;</div>
<div class="line"><a name="l00784"></a><span class="lineno">  784</span>&#160;    };</div>
<div class="line"><a name="l00785"></a><span class="lineno">  785</span>&#160; </div>
<div class="line"><a name="l00786"></a><span class="lineno">  786</span>&#160;    <span class="comment">// ThreadLocalBlocks initialization callables.</span></div>
<div class="line"><a name="l00787"></a><span class="lineno">  787</span>&#160;    <span class="keyword">using</span> ThreadLocalLhsInit =</div>
<div class="line"><a name="l00788"></a><span class="lineno">  788</span>&#160;        ThreadLocalBlocksInitialize&lt;LhsBlock, <span class="comment">/*is_rhs=*/</span><span class="keyword">false</span>&gt;;</div>
<div class="line"><a name="l00789"></a><span class="lineno">  789</span>&#160;    <span class="keyword">using</span> ThreadLocalRhsInit =</div>
<div class="line"><a name="l00790"></a><span class="lineno">  790</span>&#160;        ThreadLocalBlocksInitialize&lt;RhsBlock, <span class="comment">/*is_rhs=*/</span><span class="keyword">true</span>&gt;;</div>
<div class="line"><a name="l00791"></a><span class="lineno">  791</span>&#160; </div>
<div class="line"><a name="l00792"></a><span class="lineno">  792</span>&#160;    <span class="comment">// ThreadLocalBlocks release callables.</span></div>
<div class="line"><a name="l00793"></a><span class="lineno">  793</span>&#160;    <span class="keyword">using</span> ThreadLocalLhsRelease = ThreadLocalBlocksRelease&lt;LhsBlock&gt;;</div>
<div class="line"><a name="l00794"></a><span class="lineno">  794</span>&#160;    <span class="keyword">using</span> ThreadLocalRhsRelease = ThreadLocalBlocksRelease&lt;RhsBlock&gt;;</div>
<div class="line"><a name="l00795"></a><span class="lineno">  795</span>&#160; </div>
<div class="line"><a name="l00796"></a><span class="lineno">  796</span>&#160;    <span class="comment">// Thread local containers for Lhs/Rhs block packs. In practice only one of</span></div>
<div class="line"><a name="l00797"></a><span class="lineno">  797</span>&#160;    <span class="comment">// them will be used, depending on the shard_by_col value.</span></div>
<div class="line"><a name="l00798"></a><span class="lineno">  798</span>&#160;    Eigen::ThreadLocal&lt;ThreadLocalBlocks&lt;LhsBlock&gt;, ThreadLocalLhsInit,</div>
<div class="line"><a name="l00799"></a><span class="lineno">  799</span>&#160;                       ThreadLocalLhsRelease&gt;</div>
<div class="line"><a name="l00800"></a><span class="lineno">  800</span>&#160;        lhs_thread_local_blocks_;</div>
<div class="line"><a name="l00801"></a><span class="lineno">  801</span>&#160;    Eigen::ThreadLocal&lt;ThreadLocalBlocks&lt;RhsBlock&gt;, ThreadLocalRhsInit,</div>
<div class="line"><a name="l00802"></a><span class="lineno">  802</span>&#160;                       ThreadLocalRhsRelease&gt;</div>
<div class="line"><a name="l00803"></a><span class="lineno">  803</span>&#160;        rhs_thread_local_blocks_;</div>
<div class="line"><a name="l00804"></a><span class="lineno">  804</span>&#160; </div>
<div class="line"><a name="l00805"></a><span class="lineno">  805</span>&#160;    <span class="comment">// After a particular shard for Kth slice missed thread local execution</span></div>
<div class="line"><a name="l00806"></a><span class="lineno">  806</span>&#160;    <span class="comment">// opportunity (K-1 slice didn&#39;t complete kernels execution), we can no</span></div>
<div class="line"><a name="l00807"></a><span class="lineno">  807</span>&#160;    <span class="comment">// longer schedule K+1 and following slices in thread local mode, because</span></div>
<div class="line"><a name="l00808"></a><span class="lineno">  808</span>&#160;    <span class="comment">// there is no more guarantee that previous kernels were executed</span></div>
<div class="line"><a name="l00809"></a><span class="lineno">  809</span>&#160;    <span class="comment">// sequentially in the same thread (size is nn_ or nm_).</span></div>
<div class="line"><a name="l00810"></a><span class="lineno">  810</span>&#160;    std::atomic&lt;bool&gt;* can_use_thread_local_packed_;</div>
<div class="line"><a name="l00811"></a><span class="lineno">  811</span>&#160; </div>
<div class="line"><a name="l00812"></a><span class="lineno">  812</span>&#160;    std::atomic&lt;uint8_t&gt;** state_kernel_[P];</div>
<div class="line"><a name="l00813"></a><span class="lineno">  813</span>&#160;    <span class="comment">// state_switch_ is frequently modified by worker threads, while other</span></div>
<div class="line"><a name="l00814"></a><span class="lineno">  814</span>&#160;    <span class="comment">// fields are read-only after constructor. Let&#39;s move it to a separate cache</span></div>
<div class="line"><a name="l00815"></a><span class="lineno">  815</span>&#160;    <span class="comment">// line to reduce cache-coherency traffic.</span></div>
<div class="line"><a name="l00816"></a><span class="lineno">  816</span>&#160;    <span class="keywordtype">char</span> pad_[128];</div>
<div class="line"><a name="l00817"></a><span class="lineno">  817</span>&#160;    std::atomic&lt;Index&gt; state_packing_ready_[P];</div>
<div class="line"><a name="l00818"></a><span class="lineno">  818</span>&#160;    std::atomic&lt;Index&gt; state_switch_[P];</div>
<div class="line"><a name="l00819"></a><span class="lineno">  819</span>&#160; </div>
<div class="line"><a name="l00820"></a><span class="lineno">  820</span>&#160;    LhsBlock&amp; packed_lhs(Index m, Index k, Index m1, <span class="keywordtype">bool</span> use_thread_local) {</div>
<div class="line"><a name="l00821"></a><span class="lineno">  821</span>&#160;      <span class="keywordflow">if</span> (use_thread_local) {</div>
<div class="line"><a name="l00822"></a><span class="lineno">  822</span>&#160;        eigen_assert(!shard_by_col_);</div>
<div class="line"><a name="l00823"></a><span class="lineno">  823</span>&#160;        ThreadLocalBlocks&lt;LhsBlock&gt;&amp; blocks = lhs_thread_local_blocks_.local();</div>
<div class="line"><a name="l00824"></a><span class="lineno">  824</span>&#160; </div>
<div class="line"><a name="l00825"></a><span class="lineno">  825</span>&#160;        <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> grain_index = m1 - m * gm_;</div>
<div class="line"><a name="l00826"></a><span class="lineno">  826</span>&#160;        <span class="keywordflow">return</span> blocks.block(internal::convert_index&lt;int&gt;(grain_index)); <span class="comment">// FIXME better make ThreadLocalBlocks use Eigen::Index?</span></div>
<div class="line"><a name="l00827"></a><span class="lineno">  827</span>&#160;      } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00828"></a><span class="lineno">  828</span>&#160;        <span class="keywordflow">return</span> packed_lhs_[k % (P - 1)][m1];</div>
<div class="line"><a name="l00829"></a><span class="lineno">  829</span>&#160;      }</div>
<div class="line"><a name="l00830"></a><span class="lineno">  830</span>&#160;    }</div>
<div class="line"><a name="l00831"></a><span class="lineno">  831</span>&#160; </div>
<div class="line"><a name="l00832"></a><span class="lineno">  832</span>&#160;    RhsBlock&amp; packed_rhs(Index n, Index k, Index n1, <span class="keywordtype">bool</span> use_thread_local) {</div>
<div class="line"><a name="l00833"></a><span class="lineno">  833</span>&#160;      <span class="keywordflow">if</span> (use_thread_local) {</div>
<div class="line"><a name="l00834"></a><span class="lineno">  834</span>&#160;        eigen_assert(shard_by_col_);</div>
<div class="line"><a name="l00835"></a><span class="lineno">  835</span>&#160;        ThreadLocalBlocks&lt;RhsBlock&gt;&amp; blocks = rhs_thread_local_blocks_.local();</div>
<div class="line"><a name="l00836"></a><span class="lineno">  836</span>&#160; </div>
<div class="line"><a name="l00837"></a><span class="lineno">  837</span>&#160;        <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> grain_index = n1 - n * gn_;</div>
<div class="line"><a name="l00838"></a><span class="lineno">  838</span>&#160;        <span class="keywordflow">return</span> blocks.block(internal::convert_index&lt;int&gt;(grain_index)); <span class="comment">// FIXME better make ThreadLocalBlocks use Eigen::Index?</span></div>
<div class="line"><a name="l00839"></a><span class="lineno">  839</span>&#160;      } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00840"></a><span class="lineno">  840</span>&#160;        <span class="keywordflow">return</span> packed_rhs_[k % (P - 1)][n1];</div>
<div class="line"><a name="l00841"></a><span class="lineno">  841</span>&#160;      }</div>
<div class="line"><a name="l00842"></a><span class="lineno">  842</span>&#160;    }</div>
<div class="line"><a name="l00843"></a><span class="lineno">  843</span>&#160; </div>
<div class="line"><a name="l00844"></a><span class="lineno">  844</span>&#160;    <span class="comment">// In following two methods (pack_lhs and pack_rhs), if we know for sure</span></div>
<div class="line"><a name="l00845"></a><span class="lineno">  845</span>&#160;    <span class="comment">// that we&#39;ll be able to immediately call a kernel with packed data, and do</span></div>
<div class="line"><a name="l00846"></a><span class="lineno">  846</span>&#160;    <span class="comment">// not submit it to the thread pool, we can use thread local memory for</span></div>
<div class="line"><a name="l00847"></a><span class="lineno">  847</span>&#160;    <span class="comment">// packed data.</span></div>
<div class="line"><a name="l00848"></a><span class="lineno">  848</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l00849"></a><span class="lineno">  849</span>&#160;    <span class="comment">// We can only reliably check it if we are running all kernels in sync mode</span></div>
<div class="line"><a name="l00850"></a><span class="lineno">  850</span>&#160;    <span class="comment">// (parallelize only by sharding dim). If kernel for m==0 (n==0) is ready to</span></div>
<div class="line"><a name="l00851"></a><span class="lineno">  851</span>&#160;    <span class="comment">// run, it&#39;s guaranteed that all kernels with larger values of m (n) are</span></div>
<div class="line"><a name="l00852"></a><span class="lineno">  852</span>&#160;    <span class="comment">// also ready, because we execute them in the same order for all K slices.</span></div>
<div class="line"><a name="l00853"></a><span class="lineno">  853</span>&#160; </div>
<div class="line"><a name="l00854"></a><span class="lineno">  854</span>&#160;    <span class="keywordtype">void</span> pack_lhs(Index m, Index k) {</div>
<div class="line"><a name="l00855"></a><span class="lineno">  855</span>&#160;      <span class="keywordtype">bool</span> use_thread_local = <span class="keyword">false</span>;</div>
<div class="line"><a name="l00856"></a><span class="lineno">  856</span>&#160; </div>
<div class="line"><a name="l00857"></a><span class="lineno">  857</span>&#160;      <span class="keywordflow">if</span> (parallelize_by_sharding_dim_only_ &amp;&amp; !shard_by_col_ &amp;&amp;</div>
<div class="line"><a name="l00858"></a><span class="lineno">  858</span>&#160;          can_use_thread_local_packed_[m].load(std::memory_order_relaxed)) {</div>
<div class="line"><a name="l00859"></a><span class="lineno">  859</span>&#160;        <span class="keywordflow">if</span> (state_kernel_[k % P][m][0].load(std::memory_order_relaxed) == 1) {</div>
<div class="line"><a name="l00860"></a><span class="lineno">  860</span>&#160;          use_thread_local = <span class="keyword">true</span>;</div>
<div class="line"><a name="l00861"></a><span class="lineno">  861</span>&#160;        } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00862"></a><span class="lineno">  862</span>&#160;          <span class="comment">// If we can&#39;t guarantee that all kernels in `k` slice will be</span></div>
<div class="line"><a name="l00863"></a><span class="lineno">  863</span>&#160;          <span class="comment">// executed sequentially in current thread, it&#39;s no longer safe to use</span></div>
<div class="line"><a name="l00864"></a><span class="lineno">  864</span>&#160;          <span class="comment">// thread local memory in following slices along the k dimensions.</span></div>
<div class="line"><a name="l00865"></a><span class="lineno">  865</span>&#160;          eigen_assert(k &gt; 0);</div>
<div class="line"><a name="l00866"></a><span class="lineno">  866</span>&#160;          can_use_thread_local_packed_[m].store(<span class="keyword">false</span>,</div>
<div class="line"><a name="l00867"></a><span class="lineno">  867</span>&#160;                                                std::memory_order_relaxed);</div>
<div class="line"><a name="l00868"></a><span class="lineno">  868</span>&#160;        }</div>
<div class="line"><a name="l00869"></a><span class="lineno">  869</span>&#160;      }</div>
<div class="line"><a name="l00870"></a><span class="lineno">  870</span>&#160; </div>
<div class="line"><a name="l00871"></a><span class="lineno">  871</span>&#160;      <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> mend = m * gm_ + gm(m);</div>
<div class="line"><a name="l00872"></a><span class="lineno">  872</span>&#160;      <span class="keywordflow">for</span> (Index m1 = m * gm_; m1 &lt; mend; m1++)</div>
<div class="line"><a name="l00873"></a><span class="lineno">  873</span>&#160;        kernel_.packLhs(&amp;packed_lhs(m, k, m1, use_thread_local),</div>
<div class="line"><a name="l00874"></a><span class="lineno">  874</span>&#160;                        lhs_.getSubMapper(m1 * bm_, k * bk_), bk(k), bm(m1));</div>
<div class="line"><a name="l00875"></a><span class="lineno">  875</span>&#160; </div>
<div class="line"><a name="l00876"></a><span class="lineno">  876</span>&#160;      <span class="keywordflow">if</span> (!parallel_pack_ &amp;&amp; shard_by_col_) {</div>
<div class="line"><a name="l00877"></a><span class="lineno">  877</span>&#160;        assert(!use_thread_local);</div>
<div class="line"><a name="l00878"></a><span class="lineno">  878</span>&#160;        signal_packing(k);</div>
<div class="line"><a name="l00879"></a><span class="lineno">  879</span>&#160;      } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00880"></a><span class="lineno">  880</span>&#160;        signal_switch(k + 1);</div>
<div class="line"><a name="l00881"></a><span class="lineno">  881</span>&#160;        <span class="keywordflow">for</span> (Index n = nn_ - 1; n &gt;= 0; n--) {</div>
<div class="line"><a name="l00882"></a><span class="lineno">  882</span>&#160;          <span class="keywordtype">bool</span> sync = parallelize_by_sharding_dim_only_ || n == 0;</div>
<div class="line"><a name="l00883"></a><span class="lineno">  883</span>&#160;          signal_kernel(m, n, k, sync, use_thread_local);</div>
<div class="line"><a name="l00884"></a><span class="lineno">  884</span>&#160;        }</div>
<div class="line"><a name="l00885"></a><span class="lineno">  885</span>&#160;      }</div>
<div class="line"><a name="l00886"></a><span class="lineno">  886</span>&#160;    }</div>
<div class="line"><a name="l00887"></a><span class="lineno">  887</span>&#160; </div>
<div class="line"><a name="l00888"></a><span class="lineno">  888</span>&#160;    <span class="keywordtype">void</span> pack_rhs(Index n, Index k) {</div>
<div class="line"><a name="l00889"></a><span class="lineno">  889</span>&#160;      <span class="keywordtype">bool</span> use_thread_local = <span class="keyword">false</span>;</div>
<div class="line"><a name="l00890"></a><span class="lineno">  890</span>&#160; </div>
<div class="line"><a name="l00891"></a><span class="lineno">  891</span>&#160;      <span class="keywordflow">if</span> (parallelize_by_sharding_dim_only_ &amp;&amp; shard_by_col_ &amp;&amp;</div>
<div class="line"><a name="l00892"></a><span class="lineno">  892</span>&#160;          can_use_thread_local_packed_[n].load(std::memory_order_relaxed)) {</div>
<div class="line"><a name="l00893"></a><span class="lineno">  893</span>&#160;        <span class="keywordflow">if</span> (state_kernel_[k % P][0][n].load(std::memory_order_relaxed) == 1) {</div>
<div class="line"><a name="l00894"></a><span class="lineno">  894</span>&#160;          use_thread_local = <span class="keyword">true</span>;</div>
<div class="line"><a name="l00895"></a><span class="lineno">  895</span>&#160;        } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00896"></a><span class="lineno">  896</span>&#160;          <span class="comment">// If we can&#39;t guarantee that all kernels in `k` slice will be</span></div>
<div class="line"><a name="l00897"></a><span class="lineno">  897</span>&#160;          <span class="comment">// executed sequentially in current thread, it&#39;s no longer safe to use</span></div>
<div class="line"><a name="l00898"></a><span class="lineno">  898</span>&#160;          <span class="comment">// thread local memory in following slices along the k dimensions.</span></div>
<div class="line"><a name="l00899"></a><span class="lineno">  899</span>&#160;          eigen_assert(k &gt; 0);</div>
<div class="line"><a name="l00900"></a><span class="lineno">  900</span>&#160;          can_use_thread_local_packed_[n].store(<span class="keyword">false</span>,</div>
<div class="line"><a name="l00901"></a><span class="lineno">  901</span>&#160;                                                std::memory_order_relaxed);</div>
<div class="line"><a name="l00902"></a><span class="lineno">  902</span>&#160;        }</div>
<div class="line"><a name="l00903"></a><span class="lineno">  903</span>&#160;      }</div>
<div class="line"><a name="l00904"></a><span class="lineno">  904</span>&#160; </div>
<div class="line"><a name="l00905"></a><span class="lineno">  905</span>&#160;      <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nend = n * gn_ + gn(n);</div>
<div class="line"><a name="l00906"></a><span class="lineno">  906</span>&#160;      <span class="keywordflow">for</span> (Index n1 = n * gn_; n1 &lt; nend; n1++) {</div>
<div class="line"><a name="l00907"></a><span class="lineno">  907</span>&#160;        <span class="keywordflow">if</span> (!TensorContractionKernel::HasBeta &amp;&amp; k == 0) {</div>
<div class="line"><a name="l00908"></a><span class="lineno">  908</span>&#160;          <span class="comment">// Zero the output memory in parallel, only if contraction kernel does</span></div>
<div class="line"><a name="l00909"></a><span class="lineno">  909</span>&#160;          <span class="comment">// not support `beta`. Otherwise we will pass beta 0.0 to the first</span></div>
<div class="line"><a name="l00910"></a><span class="lineno">  910</span>&#160;          <span class="comment">// call to the `TensorContractionKernel::invoke()`.</span></div>
<div class="line"><a name="l00911"></a><span class="lineno">  911</span>&#160;          <span class="comment">//</span></div>
<div class="line"><a name="l00912"></a><span class="lineno">  912</span>&#160;          <span class="comment">// On 10000x2x10000 mm zeroing can easily take half of time. Zero (bn</span></div>
<div class="line"><a name="l00913"></a><span class="lineno">  913</span>&#160;          <span class="comment">// x m) row. Safe to do here because all kernels that will write to</span></div>
<div class="line"><a name="l00914"></a><span class="lineno">  914</span>&#160;          <span class="comment">// this memory depend on completion of this task. Note: don&#39;t call</span></div>
<div class="line"><a name="l00915"></a><span class="lineno">  915</span>&#160;          <span class="comment">// device_.fill() here. device_.fill() blocks on thread pool</span></div>
<div class="line"><a name="l00916"></a><span class="lineno">  916</span>&#160;          <span class="comment">// worker thread, which can lead to underutilization and deadlocks.</span></div>
<div class="line"><a name="l00917"></a><span class="lineno">  917</span>&#160;          std::fill_n(buffer_ + n1 * bn_ * m_, bn(n1) * m_, Scalar(0));</div>
<div class="line"><a name="l00918"></a><span class="lineno">  918</span>&#160;        }</div>
<div class="line"><a name="l00919"></a><span class="lineno">  919</span>&#160;        kernel_.packRhs(&amp;packed_rhs(n, k, n1, use_thread_local),</div>
<div class="line"><a name="l00920"></a><span class="lineno">  920</span>&#160;                        rhs_.getSubMapper(k * bk_, n1 * bn_), bk(k), bn(n1));</div>
<div class="line"><a name="l00921"></a><span class="lineno">  921</span>&#160;      }</div>
<div class="line"><a name="l00922"></a><span class="lineno">  922</span>&#160; </div>
<div class="line"><a name="l00923"></a><span class="lineno">  923</span>&#160;      <span class="keywordflow">if</span> (parallel_pack_ || shard_by_col_) {</div>
<div class="line"><a name="l00924"></a><span class="lineno">  924</span>&#160;        signal_switch(k + 1);</div>
<div class="line"><a name="l00925"></a><span class="lineno">  925</span>&#160;        <span class="keywordflow">for</span> (Index m = nm_ - 1; m &gt;= 0; m--) {</div>
<div class="line"><a name="l00926"></a><span class="lineno">  926</span>&#160;          <span class="keywordtype">bool</span> sync = parallelize_by_sharding_dim_only_ || m == 0;</div>
<div class="line"><a name="l00927"></a><span class="lineno">  927</span>&#160;          signal_kernel(m, n, k, sync, use_thread_local);</div>
<div class="line"><a name="l00928"></a><span class="lineno">  928</span>&#160;        }</div>
<div class="line"><a name="l00929"></a><span class="lineno">  929</span>&#160;      } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00930"></a><span class="lineno">  930</span>&#160;        assert(!use_thread_local);</div>
<div class="line"><a name="l00931"></a><span class="lineno">  931</span>&#160;        signal_packing(k);</div>
<div class="line"><a name="l00932"></a><span class="lineno">  932</span>&#160;      }</div>
<div class="line"><a name="l00933"></a><span class="lineno">  933</span>&#160;    }</div>
<div class="line"><a name="l00934"></a><span class="lineno">  934</span>&#160; </div>
<div class="line"><a name="l00935"></a><span class="lineno">  935</span>&#160;    <span class="keywordtype">void</span> kernel(Index m, Index n, Index k, <span class="keywordtype">bool</span> use_thread_local) {</div>
<div class="line"><a name="l00936"></a><span class="lineno">  936</span>&#160;      <span class="comment">// Note: order of iteration matters here. Iteration over m is innermost</span></div>
<div class="line"><a name="l00937"></a><span class="lineno">  937</span>&#160;      <span class="comment">// because we want to reuse the same packed rhs in consecutive tasks</span></div>
<div class="line"><a name="l00938"></a><span class="lineno">  938</span>&#160;      <span class="comment">// (rhs fits into L2$ while lhs only into L3$).</span></div>
<div class="line"><a name="l00939"></a><span class="lineno">  939</span>&#160;      <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nend = n * gn_ + gn(n);</div>
<div class="line"><a name="l00940"></a><span class="lineno">  940</span>&#160;      <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> mend = m * gm_ + gm(m);</div>
<div class="line"><a name="l00941"></a><span class="lineno">  941</span>&#160; </div>
<div class="line"><a name="l00942"></a><span class="lineno">  942</span>&#160;      <span class="comment">// NOTE: output = alpha * LHS * RHS + beta * output.</span></div>
<div class="line"><a name="l00943"></a><span class="lineno">  943</span>&#160;      <span class="keyword">const</span> Scalar alpha = Scalar(1);</div>
<div class="line"><a name="l00944"></a><span class="lineno">  944</span>&#160;      <span class="keyword">const</span> Scalar beta =</div>
<div class="line"><a name="l00945"></a><span class="lineno">  945</span>&#160;          (TensorContractionKernel::HasBeta &amp;&amp; k == 0) ? Scalar(0) : Scalar(1);</div>
<div class="line"><a name="l00946"></a><span class="lineno">  946</span>&#160; </div>
<div class="line"><a name="l00947"></a><span class="lineno">  947</span>&#160;      <span class="keywordflow">if</span> (shard_by_col_) {</div>
<div class="line"><a name="l00948"></a><span class="lineno">  948</span>&#160;        <span class="keywordflow">for</span> (Index n1 = n * gn_; n1 &lt; nend; n1++) {</div>
<div class="line"><a name="l00949"></a><span class="lineno">  949</span>&#160;          <span class="keywordflow">for</span> (Index m1 = m * gm_; m1 &lt; mend; m1++) {</div>
<div class="line"><a name="l00950"></a><span class="lineno">  950</span>&#160;            <span class="keyword">const</span> <span class="keyword">auto</span> output_mapper = output_.getSubMapper(m1 * bm_, n1 * bn_);</div>
<div class="line"><a name="l00951"></a><span class="lineno">  951</span>&#160;            kernel_.invoke(</div>
<div class="line"><a name="l00952"></a><span class="lineno">  952</span>&#160;                output_mapper,</div>
<div class="line"><a name="l00953"></a><span class="lineno">  953</span>&#160;                packed_lhs(m, k, m1, !shard_by_col_ &amp;&amp; use_thread_local),</div>
<div class="line"><a name="l00954"></a><span class="lineno">  954</span>&#160;                packed_rhs(n, k, n1, shard_by_col_ &amp;&amp; use_thread_local), bm(m1),</div>
<div class="line"><a name="l00955"></a><span class="lineno">  955</span>&#160;                bk(k), bn(n1), alpha, beta);</div>
<div class="line"><a name="l00956"></a><span class="lineno">  956</span>&#160; </div>
<div class="line"><a name="l00957"></a><span class="lineno">  957</span>&#160;            <span class="comment">// We are done with the last task for the [m1, n1] block.</span></div>
<div class="line"><a name="l00958"></a><span class="lineno">  958</span>&#160;            <span class="keywordflow">if</span> (k + 1 == nk_) {</div>
<div class="line"><a name="l00959"></a><span class="lineno">  959</span>&#160;              output_kernel_(output_mapper, tensor_contraction_params_,</div>
<div class="line"><a name="l00960"></a><span class="lineno">  960</span>&#160;                             m1 * bm_, n1 * bn_, bm(m1), bn(n1));</div>
<div class="line"><a name="l00961"></a><span class="lineno">  961</span>&#160;            }</div>
<div class="line"><a name="l00962"></a><span class="lineno">  962</span>&#160;          }</div>
<div class="line"><a name="l00963"></a><span class="lineno">  963</span>&#160;        }</div>
<div class="line"><a name="l00964"></a><span class="lineno">  964</span>&#160;      } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l00965"></a><span class="lineno">  965</span>&#160;        <span class="keywordflow">for</span> (Index m1 = m * gm_; m1 &lt; mend; m1++)</div>
<div class="line"><a name="l00966"></a><span class="lineno">  966</span>&#160;          <span class="keywordflow">for</span> (Index n1 = n * gn_; n1 &lt; nend; n1++) {</div>
<div class="line"><a name="l00967"></a><span class="lineno">  967</span>&#160;            <span class="keyword">const</span> <span class="keyword">auto</span> output_mapper = output_.getSubMapper(m1 * bm_, n1 * bn_);</div>
<div class="line"><a name="l00968"></a><span class="lineno">  968</span>&#160;            kernel_.invoke(</div>
<div class="line"><a name="l00969"></a><span class="lineno">  969</span>&#160;                output_mapper,</div>
<div class="line"><a name="l00970"></a><span class="lineno">  970</span>&#160;                packed_lhs(m, k, m1, !shard_by_col_ &amp;&amp; use_thread_local),</div>
<div class="line"><a name="l00971"></a><span class="lineno">  971</span>&#160;                packed_rhs(n, k, n1, shard_by_col_ &amp;&amp; use_thread_local), bm(m1),</div>
<div class="line"><a name="l00972"></a><span class="lineno">  972</span>&#160;                bk(k), bn(n1), alpha, beta);</div>
<div class="line"><a name="l00973"></a><span class="lineno">  973</span>&#160; </div>
<div class="line"><a name="l00974"></a><span class="lineno">  974</span>&#160;            <span class="comment">// We are done with the last task for the [m1, n1] block.</span></div>
<div class="line"><a name="l00975"></a><span class="lineno">  975</span>&#160;            <span class="keywordflow">if</span> (k + 1 == nk_) {</div>
<div class="line"><a name="l00976"></a><span class="lineno">  976</span>&#160;              output_kernel_(output_mapper, tensor_contraction_params_,</div>
<div class="line"><a name="l00977"></a><span class="lineno">  977</span>&#160;                             m1 * bm_, n1 * bn_, bm(m1), bn(n1));</div>
<div class="line"><a name="l00978"></a><span class="lineno">  978</span>&#160;            }</div>
<div class="line"><a name="l00979"></a><span class="lineno">  979</span>&#160;          }</div>
<div class="line"><a name="l00980"></a><span class="lineno">  980</span>&#160;      }</div>
<div class="line"><a name="l00981"></a><span class="lineno">  981</span>&#160;      signal_kernel(m, n, k + 1, <span class="comment">/*sync=*/</span><span class="keyword">false</span>, <span class="comment">/*use_thread_local=*/</span><span class="keyword">false</span>);</div>
<div class="line"><a name="l00982"></a><span class="lineno">  982</span>&#160;      signal_switch(k + 2);</div>
<div class="line"><a name="l00983"></a><span class="lineno">  983</span>&#160;    }</div>
<div class="line"><a name="l00984"></a><span class="lineno">  984</span>&#160; </div>
<div class="line"><a name="l00985"></a><span class="lineno">  985</span>&#160;    <span class="keywordtype">void</span> signal_packing(Index k) {</div>
<div class="line"><a name="l00986"></a><span class="lineno">  986</span>&#160;      eigen_assert(!parallel_pack_);</div>
<div class="line"><a name="l00987"></a><span class="lineno">  987</span>&#160;      <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> s = state_packing_ready_[k % P].fetch_sub(1);</div>
<div class="line"><a name="l00988"></a><span class="lineno">  988</span>&#160;      eigen_assert(s &gt; 0);</div>
<div class="line"><a name="l00989"></a><span class="lineno">  989</span>&#160;      <span class="keywordflow">if</span> (s != 1) <span class="keywordflow">return</span>;</div>
<div class="line"><a name="l00990"></a><span class="lineno">  990</span>&#160;      state_packing_ready_[k % P] = shard_by_col_ ? nm_ : nn_;</div>
<div class="line"><a name="l00991"></a><span class="lineno">  991</span>&#160;      enqueue_packing(k, shard_by_col_);</div>
<div class="line"><a name="l00992"></a><span class="lineno">  992</span>&#160;    }</div>
<div class="line"><a name="l00993"></a><span class="lineno">  993</span>&#160; </div>
<div class="line"><a name="l00994"></a><span class="lineno">  994</span>&#160;    <span class="keywordtype">void</span> signal_kernel(Index m, Index n, Index k, <span class="keywordtype">bool</span> sync,</div>
<div class="line"><a name="l00995"></a><span class="lineno">  995</span>&#160;                       <span class="keywordtype">bool</span> use_thread_local) {</div>
<div class="line"><a name="l00996"></a><span class="lineno">  996</span>&#160;      std::atomic&lt;uint8_t&gt;* state = &amp;state_kernel_[k % P][m][n];</div>
<div class="line"><a name="l00997"></a><span class="lineno">  997</span>&#160;      <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> s = state-&gt;load();</div>
<div class="line"><a name="l00998"></a><span class="lineno">  998</span>&#160;      eigen_assert(s &gt; 0);</div>
<div class="line"><a name="l00999"></a><span class="lineno">  999</span>&#160;      <span class="keywordflow">if</span> (s != 1 &amp;&amp; state-&gt;fetch_sub(1) != 1) {</div>
<div class="line"><a name="l01000"></a><span class="lineno"> 1000</span>&#160;        eigen_assert(!use_thread_local);</div>
<div class="line"><a name="l01001"></a><span class="lineno"> 1001</span>&#160;        <span class="keywordflow">return</span>;</div>
<div class="line"><a name="l01002"></a><span class="lineno"> 1002</span>&#160;      }</div>
<div class="line"><a name="l01003"></a><span class="lineno"> 1003</span>&#160;      state-&gt;store(parallel_pack_ ? 3 : 2, std::memory_order_relaxed);</div>
<div class="line"><a name="l01004"></a><span class="lineno"> 1004</span>&#160;      <span class="keywordflow">if</span> (sync) {</div>
<div class="line"><a name="l01005"></a><span class="lineno"> 1005</span>&#160;        kernel(m, n, k, use_thread_local);</div>
<div class="line"><a name="l01006"></a><span class="lineno"> 1006</span>&#160;      } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l01007"></a><span class="lineno"> 1007</span>&#160;        eigen_assert(!use_thread_local);</div>
<div class="line"><a name="l01008"></a><span class="lineno"> 1008</span>&#160;        device_.enqueueNoNotification(</div>
<div class="line"><a name="l01009"></a><span class="lineno"> 1009</span>&#160;            [=]() { kernel(m, n, k, use_thread_local); });</div>
<div class="line"><a name="l01010"></a><span class="lineno"> 1010</span>&#160;      }</div>
<div class="line"><a name="l01011"></a><span class="lineno"> 1011</span>&#160;    }</div>
<div class="line"><a name="l01012"></a><span class="lineno"> 1012</span>&#160; </div>
<div class="line"><a name="l01013"></a><span class="lineno"> 1013</span>&#160;    <span class="keywordtype">void</span> signal_switch(Index k, Index v = 1) {</div>
<div class="line"><a name="l01014"></a><span class="lineno"> 1014</span>&#160;      <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> s = state_switch_[k % P].fetch_sub(v);</div>
<div class="line"><a name="l01015"></a><span class="lineno"> 1015</span>&#160;      eigen_assert(s &gt;= v);</div>
<div class="line"><a name="l01016"></a><span class="lineno"> 1016</span>&#160;      <span class="keywordflow">if</span> (s != v) <span class="keywordflow">return</span>;</div>
<div class="line"><a name="l01017"></a><span class="lineno"> 1017</span>&#160; </div>
<div class="line"><a name="l01018"></a><span class="lineno"> 1018</span>&#160;      <span class="comment">// Ready to switch to the next k slice.</span></div>
<div class="line"><a name="l01019"></a><span class="lineno"> 1019</span>&#160;      <span class="comment">// Reset counter for the next iteration.</span></div>
<div class="line"><a name="l01020"></a><span class="lineno"> 1020</span>&#160;      state_switch_[k % P] =</div>
<div class="line"><a name="l01021"></a><span class="lineno"> 1021</span>&#160;          (parallel_pack_ ? nm_ + nn_ : (shard_by_col_ ? nn_ : nm_)) +</div>
<div class="line"><a name="l01022"></a><span class="lineno"> 1022</span>&#160;          nm_ * nn_;</div>
<div class="line"><a name="l01023"></a><span class="lineno"> 1023</span>&#160;      <span class="keywordflow">if</span> (k &lt; nk_) {</div>
<div class="line"><a name="l01024"></a><span class="lineno"> 1024</span>&#160;        <span class="comment">// Issue lhs/rhs packing. Their completion will in turn kick off</span></div>
<div class="line"><a name="l01025"></a><span class="lineno"> 1025</span>&#160;        <span class="comment">// kernels.</span></div>
<div class="line"><a name="l01026"></a><span class="lineno"> 1026</span>&#160;        <span class="keywordflow">if</span> (parallel_pack_) {</div>
<div class="line"><a name="l01027"></a><span class="lineno"> 1027</span>&#160;          enqueue_packing(k, !shard_by_col_);</div>
<div class="line"><a name="l01028"></a><span class="lineno"> 1028</span>&#160;          enqueue_packing(k, shard_by_col_);</div>
<div class="line"><a name="l01029"></a><span class="lineno"> 1029</span>&#160;        } <span class="keywordflow">else</span> <span class="keywordflow">if</span> (shard_by_col_) {</div>
<div class="line"><a name="l01030"></a><span class="lineno"> 1030</span>&#160;          enqueue_packing(k, <span class="keyword">false</span>);</div>
<div class="line"><a name="l01031"></a><span class="lineno"> 1031</span>&#160;        } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l01032"></a><span class="lineno"> 1032</span>&#160;          enqueue_packing(k, <span class="keyword">true</span>);</div>
<div class="line"><a name="l01033"></a><span class="lineno"> 1033</span>&#160;        }</div>
<div class="line"><a name="l01034"></a><span class="lineno"> 1034</span>&#160; </div>
<div class="line"><a name="l01035"></a><span class="lineno"> 1035</span>&#160;        <span class="comment">// Termination handling.</span></div>
<div class="line"><a name="l01036"></a><span class="lineno"> 1036</span>&#160;        <span class="comment">// Because kernel completion signals k + 2 switch, we need to finish nk</span></div>
<div class="line"><a name="l01037"></a><span class="lineno"> 1037</span>&#160;        <span class="comment">// + 2 slices without issuing any tasks on nk + 1 slice. So here we</span></div>
<div class="line"><a name="l01038"></a><span class="lineno"> 1038</span>&#160;        <span class="comment">// pretend that all nk + 1 packing tasks just finish instantly; so that</span></div>
<div class="line"><a name="l01039"></a><span class="lineno"> 1039</span>&#160;        <span class="comment">// nk + 2 switch only waits for completion of nk kernels.</span></div>
<div class="line"><a name="l01040"></a><span class="lineno"> 1040</span>&#160;      } <span class="keywordflow">else</span> <span class="keywordflow">if</span> (k == nk_) {</div>
<div class="line"><a name="l01041"></a><span class="lineno"> 1041</span>&#160;        signal_switch(k + 1,</div>
<div class="line"><a name="l01042"></a><span class="lineno"> 1042</span>&#160;                      parallel_pack_ ? nm_ + nn_ : (shard_by_col_ ? nn_ : nm_));</div>
<div class="line"><a name="l01043"></a><span class="lineno"> 1043</span>&#160;      } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l01044"></a><span class="lineno"> 1044</span>&#160;        done_.Notify();</div>
<div class="line"><a name="l01045"></a><span class="lineno"> 1045</span>&#160;      }</div>
<div class="line"><a name="l01046"></a><span class="lineno"> 1046</span>&#160;    }</div>
<div class="line"><a name="l01047"></a><span class="lineno"> 1047</span>&#160; </div>
<div class="line"><a name="l01048"></a><span class="lineno"> 1048</span>&#160;    <span class="comment">// Enqueue all rhs/lhs packing for k-th slice.</span></div>
<div class="line"><a name="l01049"></a><span class="lineno"> 1049</span>&#160;    <span class="keywordtype">void</span> enqueue_packing(Index k, <span class="keywordtype">bool</span> rhs) {</div>
<div class="line"><a name="l01050"></a><span class="lineno"> 1050</span>&#160;      enqueue_packing_helper(0, rhs ? nn_ : nm_, k, rhs);</div>
<div class="line"><a name="l01051"></a><span class="lineno"> 1051</span>&#160;    }</div>
<div class="line"><a name="l01052"></a><span class="lineno"> 1052</span>&#160; </div>
<div class="line"><a name="l01053"></a><span class="lineno"> 1053</span>&#160;    <span class="keywordtype">void</span> enqueue_packing_helper(Index start, Index <a class="codeRef" href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a>, Index k, <span class="keywordtype">bool</span> rhs) {</div>
<div class="line"><a name="l01054"></a><span class="lineno"> 1054</span>&#160;      <span class="keywordflow">if</span> (<a class="codeRef" href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a> - start == 1) {</div>
<div class="line"><a name="l01055"></a><span class="lineno"> 1055</span>&#160;        <span class="keywordflow">if</span> (rhs)</div>
<div class="line"><a name="l01056"></a><span class="lineno"> 1056</span>&#160;          pack_rhs(start, k);</div>
<div class="line"><a name="l01057"></a><span class="lineno"> 1057</span>&#160;        <span class="keywordflow">else</span></div>
<div class="line"><a name="l01058"></a><span class="lineno"> 1058</span>&#160;          pack_lhs(start, k);</div>
<div class="line"><a name="l01059"></a><span class="lineno"> 1059</span>&#160;      } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l01060"></a><span class="lineno"> 1060</span>&#160;        <span class="keywordflow">while</span> (<a class="codeRef" href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a> - start &gt; 1) {</div>
<div class="line"><a name="l01061"></a><span class="lineno"> 1061</span>&#160;          <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> mid = (start + <a class="codeRef" href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a>) / 2;</div>
<div class="line"><a name="l01062"></a><span class="lineno"> 1062</span>&#160;          device_.enqueueNoNotification(</div>
<div class="line"><a name="l01063"></a><span class="lineno"> 1063</span>&#160;              [=]() { enqueue_packing_helper(mid, <a class="codeRef" href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a>, k, rhs); });</div>
<div class="line"><a name="l01064"></a><span class="lineno"> 1064</span>&#160;          <a class="codeRef" href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a> = mid;</div>
<div class="line"><a name="l01065"></a><span class="lineno"> 1065</span>&#160;        }</div>
<div class="line"><a name="l01066"></a><span class="lineno"> 1066</span>&#160; </div>
<div class="line"><a name="l01067"></a><span class="lineno"> 1067</span>&#160;        <span class="comment">// Decide if we want to run first packing task (start == 0) in</span></div>
<div class="line"><a name="l01068"></a><span class="lineno"> 1068</span>&#160;        <span class="comment">// async mode if we parallelize only by sharding dim:</span></div>
<div class="line"><a name="l01069"></a><span class="lineno"> 1069</span>&#160;        <span class="comment">// (1) pack_lhs and pack_rhs call signal_switch before completing</span></div>
<div class="line"><a name="l01070"></a><span class="lineno"> 1070</span>&#160;        <span class="comment">//     all calls to signal_kernel, which in sync mode might lead</span></div>
<div class="line"><a name="l01071"></a><span class="lineno"> 1071</span>&#160;        <span class="comment">//     to the execution of the first kernel of the k+1 slice, before</span></div>
<div class="line"><a name="l01072"></a><span class="lineno"> 1072</span>&#160;        <span class="comment">//     completing a call to the last kernel of the k slice.</span></div>
<div class="line"><a name="l01073"></a><span class="lineno"> 1073</span>&#160;        <span class="comment">// (2) all pack tasks for sharded dim must be executed in a thread</span></div>
<div class="line"><a name="l01074"></a><span class="lineno"> 1074</span>&#160;        <span class="comment">//     pool to get pre-allocated thead local buffers.</span></div>
<div class="line"><a name="l01075"></a><span class="lineno"> 1075</span>&#160;        <span class="keywordtype">bool</span> pack_async =</div>
<div class="line"><a name="l01076"></a><span class="lineno"> 1076</span>&#160;          (start == 0) &amp;&amp;</div>
<div class="line"><a name="l01077"></a><span class="lineno"> 1077</span>&#160;          (parallelize_by_sharding_dim_only_&amp;&amp; shard_by_col_ == rhs) &amp;&amp;</div>
<div class="line"><a name="l01078"></a><span class="lineno"> 1078</span>&#160;          (k &gt; 0 || std::this_thread::get_id() == created_by_thread_id_);</div>
<div class="line"><a name="l01079"></a><span class="lineno"> 1079</span>&#160; </div>
<div class="line"><a name="l01080"></a><span class="lineno"> 1080</span>&#160;        <span class="keywordflow">if</span> (pack_async) {</div>
<div class="line"><a name="l01081"></a><span class="lineno"> 1081</span>&#160;          device_.enqueueNoNotification(</div>
<div class="line"><a name="l01082"></a><span class="lineno"> 1082</span>&#160;              [=]() { enqueue_packing_helper(start, <a class="codeRef" href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a>, k, rhs); });</div>
<div class="line"><a name="l01083"></a><span class="lineno"> 1083</span>&#160;        } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l01084"></a><span class="lineno"> 1084</span>&#160;          enqueue_packing_helper(start, <a class="codeRef" href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a>, k, rhs);</div>
<div class="line"><a name="l01085"></a><span class="lineno"> 1085</span>&#160;        }</div>
<div class="line"><a name="l01086"></a><span class="lineno"> 1086</span>&#160;      }</div>
<div class="line"><a name="l01087"></a><span class="lineno"> 1087</span>&#160;    }</div>
<div class="line"><a name="l01088"></a><span class="lineno"> 1088</span>&#160; </div>
<div class="line"><a name="l01089"></a><span class="lineno"> 1089</span>&#160;    <span class="comment">// Block sizes with accounting for potentially incomplete last block.</span></div>
<div class="line"><a name="l01090"></a><span class="lineno"> 1090</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> bm(Index m)<span class="keyword"> const </span>{ <span class="keywordflow">return</span> m + 1 &lt; nm0_ ? bm_ : m_ + bm_ - bm_ * nm0_; }</div>
<div class="line"><a name="l01091"></a><span class="lineno"> 1091</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> bn(Index n)<span class="keyword"> const </span>{ <span class="keywordflow">return</span> n + 1 &lt; nn0_ ? bn_ : n_ + bn_ - bn_ * nn0_; }</div>
<div class="line"><a name="l01092"></a><span class="lineno"> 1092</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> bk(Index k)<span class="keyword"> const </span>{ <span class="keywordflow">return</span> k + 1 &lt; nk_ ? bk_ : k_ + bk_ - bk_ * nk_; }</div>
<div class="line"><a name="l01093"></a><span class="lineno"> 1093</span>&#160;    <span class="comment">// Task grain sizes accounting for potentially incomplete last task.</span></div>
<div class="line"><a name="l01094"></a><span class="lineno"> 1094</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> gm(Index m)<span class="keyword"> const </span>{ <span class="keywordflow">return</span> m + 1 &lt; nm_ ? gm_ : nm0_ + gm_ - gm_ * nm_; }</div>
<div class="line"><a name="l01095"></a><span class="lineno"> 1095</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> gn(Index n)<span class="keyword"> const </span>{ <span class="keywordflow">return</span> n + 1 &lt; nn_ ? gn_ : nn0_ + gn_ - gn_ * nn_; }</div>
<div class="line"><a name="l01096"></a><span class="lineno"> 1096</span>&#160; </div>
<div class="line"><a name="l01097"></a><span class="lineno"> 1097</span>&#160;    EvalParallelContext(<span class="keyword">const</span> EvalParallelContext&amp;) = <span class="keyword">delete</span>;</div>
<div class="line"><a name="l01098"></a><span class="lineno"> 1098</span>&#160;    <span class="keywordtype">void</span> operator=(<span class="keyword">const</span> EvalParallelContext&amp;) = <span class="keyword">delete</span>;</div>
<div class="line"><a name="l01099"></a><span class="lineno"> 1099</span>&#160;  };</div>
<div class="line"><a name="l01100"></a><span class="lineno"> 1100</span>&#160; </div>
<div class="line"><a name="l01101"></a><span class="lineno"> 1101</span>&#160;  <span class="keyword">template</span> &lt;<span class="keywordtype">bool</span> lhs_inner_dim_contiguous, <span class="keywordtype">bool</span> rhs_inner_dim_contiguous,</div>
<div class="line"><a name="l01102"></a><span class="lineno"> 1102</span>&#160;            <span class="keywordtype">bool</span> rhs_inner_dim_reordered, <span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l01103"></a><span class="lineno"> 1103</span>&#160;  <span class="keyword">using</span> SyncEvalParallelContext =</div>
<div class="line"><a name="l01104"></a><span class="lineno"> 1104</span>&#160;      EvalParallelContext&lt;NoCallback, lhs_inner_dim_contiguous,</div>
<div class="line"><a name="l01105"></a><span class="lineno"> 1105</span>&#160;                          rhs_inner_dim_contiguous, rhs_inner_dim_reordered,</div>
<div class="line"><a name="l01106"></a><span class="lineno"> 1106</span>&#160;                          Alignment&gt;;</div>
<div class="line"><a name="l01107"></a><span class="lineno"> 1107</span>&#160; </div>
<div class="line"><a name="l01108"></a><span class="lineno"> 1108</span>&#160;  <span class="comment">// ------------------------------------------------------------------------ //</span></div>
<div class="line"><a name="l01109"></a><span class="lineno"> 1109</span>&#160; </div>
<div class="line"><a name="l01110"></a><span class="lineno"> 1110</span>&#160;  <span class="comment">// EvalShardedByInnerDimContext orchestrates sync/async contraction</span></div>
<div class="line"><a name="l01111"></a><span class="lineno"> 1111</span>&#160;  <span class="comment">// evaluation, when we shard by inner dimension. When it is executed in</span></div>
<div class="line"><a name="l01112"></a><span class="lineno"> 1112</span>&#160;  <span class="comment">// asynchronous mode, it owns all the shared state that might be accessible by</span></div>
<div class="line"><a name="l01113"></a><span class="lineno"> 1113</span>&#160;  <span class="comment">// block processing tasks.</span></div>
<div class="line"><a name="l01114"></a><span class="lineno"> 1114</span>&#160; </div>
<div class="line"><a name="l01115"></a><span class="lineno"> 1115</span>&#160;  <span class="keyword">template</span> &lt;<span class="keyword">typename</span> DoneCallback&gt;</div>
<div class="line"><a name="l01116"></a><span class="lineno"> 1116</span>&#160;  <span class="keyword">struct </span>EvalShardedByInnerDimContext {</div>
<div class="line"><a name="l01117"></a><span class="lineno"> 1117</span>&#160;    EvalShardedByInnerDimContext(<span class="keyword">const</span> Self* <span class="keyword">self</span>, <span class="keywordtype">int</span> num_threads,</div>
<div class="line"><a name="l01118"></a><span class="lineno"> 1118</span>&#160;                                 Scalar* result_buffer,</div>
<div class="line"><a name="l01119"></a><span class="lineno"> 1119</span>&#160;                                 Index m_size, Index n_size, Index k_size,</div>
<div class="line"><a name="l01120"></a><span class="lineno"> 1120</span>&#160;                                 DoneCallback done_callback)</div>
<div class="line"><a name="l01121"></a><span class="lineno"> 1121</span>&#160;        : evaluator(self),</div>
<div class="line"><a name="l01122"></a><span class="lineno"> 1122</span>&#160;          m_lhs_inner_dim_contiguous(evaluator-&gt;m_lhs_inner_dim_contiguous),</div>
<div class="line"><a name="l01123"></a><span class="lineno"> 1123</span>&#160;          m_rhs_inner_dim_contiguous(evaluator-&gt;m_rhs_inner_dim_contiguous),</div>
<div class="line"><a name="l01124"></a><span class="lineno"> 1124</span>&#160;          m_rhs_inner_dim_reordered(evaluator-&gt;m_rhs_inner_dim_reordered),</div>
<div class="line"><a name="l01125"></a><span class="lineno"> 1125</span>&#160;          result(result_buffer),</div>
<div class="line"><a name="l01126"></a><span class="lineno"> 1126</span>&#160;          m(m_size),</div>
<div class="line"><a name="l01127"></a><span class="lineno"> 1127</span>&#160;          n(n_size),</div>
<div class="line"><a name="l01128"></a><span class="lineno"> 1128</span>&#160;          k(k_size),</div>
<div class="line"><a name="l01129"></a><span class="lineno"> 1129</span>&#160;          done(std::move(done_callback)),</div>
<div class="line"><a name="l01130"></a><span class="lineno"> 1130</span>&#160;          buffer_size_bytes(m * n * sizeof(Scalar)),</div>
<div class="line"><a name="l01131"></a><span class="lineno"> 1131</span>&#160;          block_size(blockSize(k, num_threads)),</div>
<div class="line"><a name="l01132"></a><span class="lineno"> 1132</span>&#160;          num_blocks(divup&lt;<a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>&gt;(k, block_size)),</div>
<div class="line"><a name="l01133"></a><span class="lineno"> 1133</span>&#160;          num_pending_blocks(internal::convert_index&lt;int&gt;(num_blocks)),</div>
<div class="line"><a name="l01134"></a><span class="lineno"> 1134</span>&#160;          l0_ranges(divup&lt;<a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a>&gt;(num_blocks, l0_size)),</div>
<div class="line"><a name="l01135"></a><span class="lineno"> 1135</span>&#160;          l0_state(l0_ranges),</div>
<div class="line"><a name="l01136"></a><span class="lineno"> 1136</span>&#160;          block_buffers(num_blocks) {</div>
<div class="line"><a name="l01137"></a><span class="lineno"> 1137</span>&#160;      <span class="comment">// Keep count of pending gemm tasks for each l0 range.</span></div>
<div class="line"><a name="l01138"></a><span class="lineno"> 1138</span>&#160;      <span class="keywordflow">for</span> (<span class="keywordtype">int</span> i = 0; i &lt; l0_ranges; ++i) {</div>
<div class="line"><a name="l01139"></a><span class="lineno"> 1139</span>&#160;        <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> num_pending_tasks = actualRangeSize(l0_ranges, l0_size, i);</div>
<div class="line"><a name="l01140"></a><span class="lineno"> 1140</span>&#160;        l0_state.emplace_back(internal::convert_index&lt;int&gt;(num_pending_tasks));</div>
<div class="line"><a name="l01141"></a><span class="lineno"> 1141</span>&#160;      }</div>
<div class="line"><a name="l01142"></a><span class="lineno"> 1142</span>&#160; </div>
<div class="line"><a name="l01143"></a><span class="lineno"> 1143</span>&#160;      <span class="comment">// Allocate temporary buffers for each block.</span></div>
<div class="line"><a name="l01144"></a><span class="lineno"> 1144</span>&#160;      <span class="keywordflow">for</span> (Index block_idx = 0; block_idx &lt; num_blocks; ++block_idx) {</div>
<div class="line"><a name="l01145"></a><span class="lineno"> 1145</span>&#160;        Scalar* buf = block_idx == 0</div>
<div class="line"><a name="l01146"></a><span class="lineno"> 1146</span>&#160;                          ? result</div>
<div class="line"><a name="l01147"></a><span class="lineno"> 1147</span>&#160;                          : <span class="keyword">static_cast&lt;</span>Scalar*<span class="keyword">&gt;</span>(evaluator-&gt;m_device.allocate(</div>
<div class="line"><a name="l01148"></a><span class="lineno"> 1148</span>&#160;                                buffer_size_bytes));</div>
<div class="line"><a name="l01149"></a><span class="lineno"> 1149</span>&#160;        block_buffers.emplace_back(buf);</div>
<div class="line"><a name="l01150"></a><span class="lineno"> 1150</span>&#160;      }</div>
<div class="line"><a name="l01151"></a><span class="lineno"> 1151</span>&#160;    }</div>
<div class="line"><a name="l01152"></a><span class="lineno"> 1152</span>&#160; </div>
<div class="line"><a name="l01153"></a><span class="lineno"> 1153</span>&#160;    ~EvalShardedByInnerDimContext() {</div>
<div class="line"><a name="l01154"></a><span class="lineno"> 1154</span>&#160;      <span class="keywordflow">for</span> (Index i = 1; i &lt; num_blocks; ++i) {</div>
<div class="line"><a name="l01155"></a><span class="lineno"> 1155</span>&#160;        evaluator-&gt;m_device.deallocate(block_buffers[i]);</div>
<div class="line"><a name="l01156"></a><span class="lineno"> 1156</span>&#160;      }</div>
<div class="line"><a name="l01157"></a><span class="lineno"> 1157</span>&#160;    }</div>
<div class="line"><a name="l01158"></a><span class="lineno"> 1158</span>&#160; </div>
<div class="line"><a name="l01159"></a><span class="lineno"> 1159</span>&#160;    <span class="keyword">template</span> &lt;<span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l01160"></a><span class="lineno"> 1160</span>&#160;    <span class="keywordtype">void</span> run() {</div>
<div class="line"><a name="l01161"></a><span class="lineno"> 1161</span>&#160;      Barrier barrier(internal::convert_index&lt;int&gt;(num_blocks));</div>
<div class="line"><a name="l01162"></a><span class="lineno"> 1162</span>&#160;      eval&lt;Alignment&gt;(barrier, 0, num_blocks);</div>
<div class="line"><a name="l01163"></a><span class="lineno"> 1163</span>&#160;      barrier.Wait();</div>
<div class="line"><a name="l01164"></a><span class="lineno"> 1164</span>&#160; </div>
<div class="line"><a name="l01165"></a><span class="lineno"> 1165</span>&#160;      <span class="comment">// Aggregate partial sums from l0 ranges.</span></div>
<div class="line"><a name="l01166"></a><span class="lineno"> 1166</span>&#160;      aggregateL0Blocks&lt;Alignment&gt;();</div>
<div class="line"><a name="l01167"></a><span class="lineno"> 1167</span>&#160; </div>
<div class="line"><a name="l01168"></a><span class="lineno"> 1168</span>&#160;      <span class="comment">// Apply output kernel.</span></div>
<div class="line"><a name="l01169"></a><span class="lineno"> 1169</span>&#160;      applyOutputKernel();</div>
<div class="line"><a name="l01170"></a><span class="lineno"> 1170</span>&#160;    }</div>
<div class="line"><a name="l01171"></a><span class="lineno"> 1171</span>&#160; </div>
<div class="line"><a name="l01172"></a><span class="lineno"> 1172</span>&#160;    <span class="keyword">template</span> &lt;<span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l01173"></a><span class="lineno"> 1173</span>&#160;    <span class="keywordtype">void</span> runAsync() {</div>
<div class="line"><a name="l01174"></a><span class="lineno"> 1174</span>&#160;      evalAsync&lt;Alignment&gt;(0, num_blocks);</div>
<div class="line"><a name="l01175"></a><span class="lineno"> 1175</span>&#160;    }</div>
<div class="line"><a name="l01176"></a><span class="lineno"> 1176</span>&#160; </div>
<div class="line"><a name="l01177"></a><span class="lineno"> 1177</span>&#160;   <span class="keyword">private</span>:</div>
<div class="line"><a name="l01178"></a><span class="lineno"> 1178</span>&#160;    <span class="comment">// The underlying GEMM kernel assumes that k is a multiple of</span></div>
<div class="line"><a name="l01179"></a><span class="lineno"> 1179</span>&#160;    <span class="comment">// the packet size and subtle breakage occurs if this is violated.</span></div>
<div class="line"><a name="l01180"></a><span class="lineno"> 1180</span>&#160;    <span class="keyword">static</span> <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> packet_size = internal::packet_traits&lt;RhsScalar&gt;::size;</div>
<div class="line"><a name="l01181"></a><span class="lineno"> 1181</span>&#160; </div>
<div class="line"><a name="l01182"></a><span class="lineno"> 1182</span>&#160;    <span class="keyword">const</span> Self* evaluator;  <span class="comment">// TensorContraction evaluator</span></div>
<div class="line"><a name="l01183"></a><span class="lineno"> 1183</span>&#160; </div>
<div class="line"><a name="l01184"></a><span class="lineno"> 1184</span>&#160;    <span class="comment">// These fields required fromTENSOR_CONTRACTION_DISPATCH macro.</span></div>
<div class="line"><a name="l01185"></a><span class="lineno"> 1185</span>&#160;    <span class="keywordtype">bool</span> m_lhs_inner_dim_contiguous;</div>
<div class="line"><a name="l01186"></a><span class="lineno"> 1186</span>&#160;    <span class="keywordtype">bool</span> m_rhs_inner_dim_contiguous;</div>
<div class="line"><a name="l01187"></a><span class="lineno"> 1187</span>&#160;    <span class="keywordtype">bool</span> m_rhs_inner_dim_reordered;</div>
<div class="line"><a name="l01188"></a><span class="lineno"> 1188</span>&#160; </div>
<div class="line"><a name="l01189"></a><span class="lineno"> 1189</span>&#160;    Scalar* result;</div>
<div class="line"><a name="l01190"></a><span class="lineno"> 1190</span>&#160; </div>
<div class="line"><a name="l01191"></a><span class="lineno"> 1191</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> m;</div>
<div class="line"><a name="l01192"></a><span class="lineno"> 1192</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> n;</div>
<div class="line"><a name="l01193"></a><span class="lineno"> 1193</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> k;</div>
<div class="line"><a name="l01194"></a><span class="lineno"> 1194</span>&#160; </div>
<div class="line"><a name="l01195"></a><span class="lineno"> 1195</span>&#160;    DoneCallback done;</div>
<div class="line"><a name="l01196"></a><span class="lineno"> 1196</span>&#160; </div>
<div class="line"><a name="l01197"></a><span class="lineno"> 1197</span>&#160;    <span class="comment">// ----------------------------------------------------------------------//</span></div>
<div class="line"><a name="l01198"></a><span class="lineno"> 1198</span>&#160;    <span class="comment">// Algorithm parameters.</span></div>
<div class="line"><a name="l01199"></a><span class="lineno"> 1199</span>&#160; </div>
<div class="line"><a name="l01200"></a><span class="lineno"> 1200</span>&#160;    <span class="comment">// We will compute partial results into the buffers of this size.</span></div>
<div class="line"><a name="l01201"></a><span class="lineno"> 1201</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> buffer_size_bytes;</div>
<div class="line"><a name="l01202"></a><span class="lineno"> 1202</span>&#160; </div>
<div class="line"><a name="l01203"></a><span class="lineno"> 1203</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> block_size;</div>
<div class="line"><a name="l01204"></a><span class="lineno"> 1204</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> num_blocks;</div>
<div class="line"><a name="l01205"></a><span class="lineno"> 1205</span>&#160; </div>
<div class="line"><a name="l01206"></a><span class="lineno"> 1206</span>&#160;    <span class="comment">// Keep track of pending tasks when evaluate in async mode.</span></div>
<div class="line"><a name="l01207"></a><span class="lineno"> 1207</span>&#160;    std::atomic&lt;int&gt; num_pending_blocks;</div>
<div class="line"><a name="l01208"></a><span class="lineno"> 1208</span>&#160; </div>
<div class="line"><a name="l01209"></a><span class="lineno"> 1209</span>&#160;    <span class="comment">// We compute partial gemm results in parallel, and to get the final result</span></div>
<div class="line"><a name="l01210"></a><span class="lineno"> 1210</span>&#160;    <span class="comment">// we need to add them all together. For the large number of threads (&gt;= 48)</span></div>
<div class="line"><a name="l01211"></a><span class="lineno"> 1211</span>&#160;    <span class="comment">// this adds a very expensive sequential step at the end.</span></div>
<div class="line"><a name="l01212"></a><span class="lineno"> 1212</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l01213"></a><span class="lineno"> 1213</span>&#160;    <span class="comment">// We split the [0, num_blocks) into small ranges, and when a task for the</span></div>
<div class="line"><a name="l01214"></a><span class="lineno"> 1214</span>&#160;    <span class="comment">// block finishes its partial gemm computation, it checks if it was the last</span></div>
<div class="line"><a name="l01215"></a><span class="lineno"> 1215</span>&#160;    <span class="comment">// gemm in the range, and if so, it will add all blocks of the range.</span></div>
<div class="line"><a name="l01216"></a><span class="lineno"> 1216</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l01217"></a><span class="lineno"> 1217</span>&#160;    <span class="comment">// After all tasks done, we need to add only these pre-aggregated blocks.</span></div>
<div class="line"><a name="l01218"></a><span class="lineno"> 1218</span>&#160; </div>
<div class="line"><a name="l01219"></a><span class="lineno"> 1219</span>&#160;    <span class="comment">// For now we use just a single level of ranges to compute pre-aggregated</span></div>
<div class="line"><a name="l01220"></a><span class="lineno"> 1220</span>&#160;    <span class="comment">// partial sums, but in general we can use more layers to compute tree</span></div>
<div class="line"><a name="l01221"></a><span class="lineno"> 1221</span>&#160;    <span class="comment">// aggregation in parallel and reduce the size of the sequential step.</span></div>
<div class="line"><a name="l01222"></a><span class="lineno"> 1222</span>&#160;    <span class="comment">//</span></div>
<div class="line"><a name="l01223"></a><span class="lineno"> 1223</span>&#160;    <span class="comment">// TODO(ezhulenev): Add multilevel tree aggregation? Probably will make</span></div>
<div class="line"><a name="l01224"></a><span class="lineno"> 1224</span>&#160;    <span class="comment">// sense only if number of threads &gt;= ~128?</span></div>
<div class="line"><a name="l01225"></a><span class="lineno"> 1225</span>&#160;    <span class="keyword">static</span> <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> l0_size = 4;</div>
<div class="line"><a name="l01226"></a><span class="lineno"> 1226</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> l0_ranges;</div>
<div class="line"><a name="l01227"></a><span class="lineno"> 1227</span>&#160; </div>
<div class="line"><a name="l01228"></a><span class="lineno"> 1228</span>&#160;    <span class="comment">// Keep count of pending gemm tasks for each l0 range.</span></div>
<div class="line"><a name="l01229"></a><span class="lineno"> 1229</span>&#160;    MaxSizeVector&lt;std::atomic&lt;int&gt;&gt; l0_state;  <span class="comment">// [0, l0_ranges)</span></div>
<div class="line"><a name="l01230"></a><span class="lineno"> 1230</span>&#160; </div>
<div class="line"><a name="l01231"></a><span class="lineno"> 1231</span>&#160;    <span class="comment">// Buffers allocated for each temporary block computation.</span></div>
<div class="line"><a name="l01232"></a><span class="lineno"> 1232</span>&#160;    MaxSizeVector&lt;Scalar*&gt; block_buffers;  <span class="comment">// [0, num_blocks)</span></div>
<div class="line"><a name="l01233"></a><span class="lineno"> 1233</span>&#160; </div>
<div class="line"><a name="l01234"></a><span class="lineno"> 1234</span>&#160;    <span class="keyword">template</span> &lt;<span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l01235"></a><span class="lineno"> 1235</span>&#160;    <span class="keywordtype">void</span> processBlock(Index block_idx, Index begin, Index <a class="codeRef" href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a>) {</div>
<div class="line"><a name="l01236"></a><span class="lineno"> 1236</span>&#160;      Scalar* buf = block_buffers[block_idx];</div>
<div class="line"><a name="l01237"></a><span class="lineno"> 1237</span>&#160; </div>
<div class="line"><a name="l01238"></a><span class="lineno"> 1238</span>&#160;      TENSOR_CONTRACTION_DISPATCH(</div>
<div class="line"><a name="l01239"></a><span class="lineno"> 1239</span>&#160;          evaluator-&gt;template evalGemmPartialWithoutOutputKernel, Alignment,</div>
<div class="line"><a name="l01240"></a><span class="lineno"> 1240</span>&#160;          (buf, begin, <a class="codeRef" href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a>,</div>
<div class="line"><a name="l01241"></a><span class="lineno"> 1241</span>&#160;           <span class="comment">/*num_threads=*/</span>internal::convert_index&lt;int&gt;(num_blocks)));</div>
<div class="line"><a name="l01242"></a><span class="lineno"> 1242</span>&#160; </div>
<div class="line"><a name="l01243"></a><span class="lineno"> 1243</span>&#160;      <span class="comment">// Check if it was the last task in l0 range.</span></div>
<div class="line"><a name="l01244"></a><span class="lineno"> 1244</span>&#160;      <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> l0_index = block_idx / l0_size;</div>
<div class="line"><a name="l01245"></a><span class="lineno"> 1245</span>&#160;      <span class="keyword">const</span> <span class="keywordtype">int</span> v = l0_state[l0_index].fetch_sub(1);</div>
<div class="line"><a name="l01246"></a><span class="lineno"> 1246</span>&#160;      eigen_assert(v &gt;= 1);</div>
<div class="line"><a name="l01247"></a><span class="lineno"> 1247</span>&#160; </div>
<div class="line"><a name="l01248"></a><span class="lineno"> 1248</span>&#160;      <span class="comment">// If we processed the last block of the range, we can aggregate all</span></div>
<div class="line"><a name="l01249"></a><span class="lineno"> 1249</span>&#160;      <span class="comment">// partial results into the first block of the range.</span></div>
<div class="line"><a name="l01250"></a><span class="lineno"> 1250</span>&#160;      <span class="keywordflow">if</span> (v == 1) {</div>
<div class="line"><a name="l01251"></a><span class="lineno"> 1251</span>&#160;        <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> rng_size = actualRangeSize(l0_ranges, l0_size, l0_index);</div>
<div class="line"><a name="l01252"></a><span class="lineno"> 1252</span>&#160;        <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> dst_block_idx = l0_index * l0_size;</div>
<div class="line"><a name="l01253"></a><span class="lineno"> 1253</span>&#160; </div>
<div class="line"><a name="l01254"></a><span class="lineno"> 1254</span>&#160;        <span class="keywordflow">if</span> (rng_size == l0_size) {</div>
<div class="line"><a name="l01255"></a><span class="lineno"> 1255</span>&#160;          addAllToBuffer&lt;Alignment&gt;(</div>
<div class="line"><a name="l01256"></a><span class="lineno"> 1256</span>&#160;              m * n,</div>
<div class="line"><a name="l01257"></a><span class="lineno"> 1257</span>&#160;              <span class="comment">/*src_buf0=*/</span>block_buffers[dst_block_idx + 1],</div>
<div class="line"><a name="l01258"></a><span class="lineno"> 1258</span>&#160;              <span class="comment">/*src_buf1=*/</span>block_buffers[dst_block_idx + 2],</div>
<div class="line"><a name="l01259"></a><span class="lineno"> 1259</span>&#160;              <span class="comment">/*src_buf2=*/</span>block_buffers[dst_block_idx + 3],</div>
<div class="line"><a name="l01260"></a><span class="lineno"> 1260</span>&#160;              <span class="comment">/*dst_buf= */</span> block_buffers[dst_block_idx]);</div>
<div class="line"><a name="l01261"></a><span class="lineno"> 1261</span>&#160;        } <span class="keywordflow">else</span> {</div>
<div class="line"><a name="l01262"></a><span class="lineno"> 1262</span>&#160;          <span class="comment">// Aggregate blocks of potentially incomplete last range.</span></div>
<div class="line"><a name="l01263"></a><span class="lineno"> 1263</span>&#160;          <span class="keywordflow">for</span> (<span class="keywordtype">int</span> i = 1; i &lt; rng_size; ++i) {</div>
<div class="line"><a name="l01264"></a><span class="lineno"> 1264</span>&#160;            addToBuffer&lt;Alignment&gt;(m * n,</div>
<div class="line"><a name="l01265"></a><span class="lineno"> 1265</span>&#160;                                   <span class="comment">/*src_buf=*/</span>block_buffers[dst_block_idx + i],</div>
<div class="line"><a name="l01266"></a><span class="lineno"> 1266</span>&#160;                                   <span class="comment">/*dst_buf=*/</span>block_buffers[dst_block_idx]);</div>
<div class="line"><a name="l01267"></a><span class="lineno"> 1267</span>&#160;          }</div>
<div class="line"><a name="l01268"></a><span class="lineno"> 1268</span>&#160;        }</div>
<div class="line"><a name="l01269"></a><span class="lineno"> 1269</span>&#160;      }</div>
<div class="line"><a name="l01270"></a><span class="lineno"> 1270</span>&#160;    }</div>
<div class="line"><a name="l01271"></a><span class="lineno"> 1271</span>&#160; </div>
<div class="line"><a name="l01272"></a><span class="lineno"> 1272</span>&#160;    <span class="comment">// Aggregate partial sums from l0 ranges.</span></div>
<div class="line"><a name="l01273"></a><span class="lineno"> 1273</span>&#160;    <span class="keyword">template</span> &lt;<span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l01274"></a><span class="lineno"> 1274</span>&#160;    <span class="keywordtype">void</span> aggregateL0Blocks()<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01275"></a><span class="lineno"> 1275</span>&#160;      <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> l0_index = 1;</div>
<div class="line"><a name="l01276"></a><span class="lineno"> 1276</span>&#160; </div>
<div class="line"><a name="l01277"></a><span class="lineno"> 1277</span>&#160;      <span class="keywordflow">for</span> (; l0_index + 2 &lt; l0_ranges; l0_index += 3) {</div>
<div class="line"><a name="l01278"></a><span class="lineno"> 1278</span>&#160;        addAllToBuffer&lt;Alignment&gt;(</div>
<div class="line"><a name="l01279"></a><span class="lineno"> 1279</span>&#160;            m * n,</div>
<div class="line"><a name="l01280"></a><span class="lineno"> 1280</span>&#160;            <span class="comment">/*src_buf0=*/</span>block_buffers[(l0_index + 0) * l0_size],</div>
<div class="line"><a name="l01281"></a><span class="lineno"> 1281</span>&#160;            <span class="comment">/*src_buf1=*/</span>block_buffers[(l0_index + 1) * l0_size],</div>
<div class="line"><a name="l01282"></a><span class="lineno"> 1282</span>&#160;            <span class="comment">/*src_buf2=*/</span>block_buffers[(l0_index + 2) * l0_size],</div>
<div class="line"><a name="l01283"></a><span class="lineno"> 1283</span>&#160;            <span class="comment">/*dst_buf= */</span> block_buffers[0]);</div>
<div class="line"><a name="l01284"></a><span class="lineno"> 1284</span>&#160;      }</div>
<div class="line"><a name="l01285"></a><span class="lineno"> 1285</span>&#160; </div>
<div class="line"><a name="l01286"></a><span class="lineno"> 1286</span>&#160;      <span class="keywordflow">for</span> (; l0_index &lt; l0_ranges; ++l0_index) {</div>
<div class="line"><a name="l01287"></a><span class="lineno"> 1287</span>&#160;        addToBuffer&lt;Alignment&gt;(m * n, block_buffers[l0_index * l0_size],</div>
<div class="line"><a name="l01288"></a><span class="lineno"> 1288</span>&#160;                               block_buffers[0]);</div>
<div class="line"><a name="l01289"></a><span class="lineno"> 1289</span>&#160;      }</div>
<div class="line"><a name="l01290"></a><span class="lineno"> 1290</span>&#160;    }</div>
<div class="line"><a name="l01291"></a><span class="lineno"> 1291</span>&#160; </div>
<div class="line"><a name="l01292"></a><span class="lineno"> 1292</span>&#160;    <span class="keywordtype">void</span> applyOutputKernel()<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01293"></a><span class="lineno"> 1293</span>&#160;      <span class="keyword">typedef</span> internal::blas_data_mapper&lt;Scalar, Index, ColMajor&gt; OutputMapper;</div>
<div class="line"><a name="l01294"></a><span class="lineno"> 1294</span>&#160;      evaluator-&gt;m_output_kernel(</div>
<div class="line"><a name="l01295"></a><span class="lineno"> 1295</span>&#160;          OutputMapper(result, m), evaluator-&gt;m_tensor_contraction_params,</div>
<div class="line"><a name="l01296"></a><span class="lineno"> 1296</span>&#160;          <span class="keyword">static_cast&lt;</span><a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Eigen::Index</a><span class="keyword">&gt;</span>(0), <span class="keyword">static_cast&lt;</span><a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Eigen::Index</a><span class="keyword">&gt;</span>(0), m, n);</div>
<div class="line"><a name="l01297"></a><span class="lineno"> 1297</span>&#160;    }</div>
<div class="line"><a name="l01298"></a><span class="lineno"> 1298</span>&#160; </div>
<div class="line"><a name="l01299"></a><span class="lineno"> 1299</span>&#160;    <span class="comment">// Compute block size with accounting for potentially incomplete last block.</span></div>
<div class="line"><a name="l01300"></a><span class="lineno"> 1300</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> actualBlockSize(Index block_idx)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01301"></a><span class="lineno"> 1301</span>&#160;      <span class="keywordflow">return</span> block_idx + 1 &lt; num_blocks</div>
<div class="line"><a name="l01302"></a><span class="lineno"> 1302</span>&#160;                 ? block_size</div>
<div class="line"><a name="l01303"></a><span class="lineno"> 1303</span>&#160;                 : k + block_size - block_size * num_blocks;</div>
<div class="line"><a name="l01304"></a><span class="lineno"> 1304</span>&#160;    };</div>
<div class="line"><a name="l01305"></a><span class="lineno"> 1305</span>&#160; </div>
<div class="line"><a name="l01306"></a><span class="lineno"> 1306</span>&#160;    <span class="comment">// Compute range size with accounting for potentially incomplete last range.</span></div>
<div class="line"><a name="l01307"></a><span class="lineno"> 1307</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> actualRangeSize(Index num_ranges, Index range_size,</div>
<div class="line"><a name="l01308"></a><span class="lineno"> 1308</span>&#160;                          Index range_idx)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01309"></a><span class="lineno"> 1309</span>&#160;      eigen_assert(range_idx &lt; num_ranges);</div>
<div class="line"><a name="l01310"></a><span class="lineno"> 1310</span>&#160;      <span class="keywordflow">return</span> range_idx + 1 &lt; num_ranges</div>
<div class="line"><a name="l01311"></a><span class="lineno"> 1311</span>&#160;                 ? range_size</div>
<div class="line"><a name="l01312"></a><span class="lineno"> 1312</span>&#160;                 : num_blocks + range_size - range_size * num_ranges;</div>
<div class="line"><a name="l01313"></a><span class="lineno"> 1313</span>&#160;    };</div>
<div class="line"><a name="l01314"></a><span class="lineno"> 1314</span>&#160; </div>
<div class="line"><a name="l01315"></a><span class="lineno"> 1315</span>&#160;    <span class="keyword">template</span> &lt;<span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l01316"></a><span class="lineno"> 1316</span>&#160;    EIGEN_STRONG_INLINE <span class="keyword">static</span> <span class="keywordtype">void</span> addToBuffer(<span class="keywordtype">size_t</span> n, <span class="keyword">const</span> Scalar* src_buf,</div>
<div class="line"><a name="l01317"></a><span class="lineno"> 1317</span>&#160;                                                Scalar* tgt_buf) {</div>
<div class="line"><a name="l01318"></a><span class="lineno"> 1318</span>&#160;      <span class="keyword">const</span> <span class="keywordtype">int</span> output_packet_size =</div>
<div class="line"><a name="l01319"></a><span class="lineno"> 1319</span>&#160;          internal::unpacket_traits&lt;PacketReturnType&gt;::size;</div>
<div class="line"><a name="l01320"></a><span class="lineno"> 1320</span>&#160;      <span class="keywordtype">size_t</span> i = 0;</div>
<div class="line"><a name="l01321"></a><span class="lineno"> 1321</span>&#160;      <span class="keyword">const</span> <span class="keywordtype">size_t</span> num_packets = n / output_packet_size;</div>
<div class="line"><a name="l01322"></a><span class="lineno"> 1322</span>&#160;      <span class="keywordflow">for</span> (; i &lt; output_packet_size * num_packets; i += output_packet_size) {</div>
<div class="line"><a name="l01323"></a><span class="lineno"> 1323</span>&#160;        <span class="keyword">const</span> PacketReturnType src_val =</div>
<div class="line"><a name="l01324"></a><span class="lineno"> 1324</span>&#160;            internal::pload&lt;PacketReturnType&gt;(src_buf + i);</div>
<div class="line"><a name="l01325"></a><span class="lineno"> 1325</span>&#160;        <span class="keyword">const</span> PacketReturnType tgt_val =</div>
<div class="line"><a name="l01326"></a><span class="lineno"> 1326</span>&#160;            internal::ploadt&lt;PacketReturnType, Alignment&gt;(tgt_buf + i);</div>
<div class="line"><a name="l01327"></a><span class="lineno"> 1327</span>&#160;        <span class="keyword">const</span> PacketReturnType sum = internal::padd(src_val, tgt_val);</div>
<div class="line"><a name="l01328"></a><span class="lineno"> 1328</span>&#160;        internal::pstoret&lt;Scalar, PacketReturnType, Alignment&gt;(tgt_buf + i,</div>
<div class="line"><a name="l01329"></a><span class="lineno"> 1329</span>&#160;                                                               sum);</div>
<div class="line"><a name="l01330"></a><span class="lineno"> 1330</span>&#160;      }</div>
<div class="line"><a name="l01331"></a><span class="lineno"> 1331</span>&#160;      <span class="keywordflow">for</span> (; i &lt; n; ++i) {</div>
<div class="line"><a name="l01332"></a><span class="lineno"> 1332</span>&#160;        tgt_buf[i] += src_buf[i];</div>
<div class="line"><a name="l01333"></a><span class="lineno"> 1333</span>&#160;      }</div>
<div class="line"><a name="l01334"></a><span class="lineno"> 1334</span>&#160;    }</div>
<div class="line"><a name="l01335"></a><span class="lineno"> 1335</span>&#160; </div>
<div class="line"><a name="l01336"></a><span class="lineno"> 1336</span>&#160;    <span class="keyword">template</span> &lt;<span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l01337"></a><span class="lineno"> 1337</span>&#160;    EIGEN_STRONG_INLINE <span class="keyword">static</span> <span class="keywordtype">void</span> addAllToBuffer(<span class="keywordtype">size_t</span> n,</div>
<div class="line"><a name="l01338"></a><span class="lineno"> 1338</span>&#160;                                                   <span class="keyword">const</span> Scalar* src_buf0,</div>
<div class="line"><a name="l01339"></a><span class="lineno"> 1339</span>&#160;                                                   <span class="keyword">const</span> Scalar* src_buf1,</div>
<div class="line"><a name="l01340"></a><span class="lineno"> 1340</span>&#160;                                                   <span class="keyword">const</span> Scalar* src_buf2,</div>
<div class="line"><a name="l01341"></a><span class="lineno"> 1341</span>&#160;                                                   Scalar* dst_buf) {</div>
<div class="line"><a name="l01342"></a><span class="lineno"> 1342</span>&#160;      using ::Eigen::internal::padd;</div>
<div class="line"><a name="l01343"></a><span class="lineno"> 1343</span>&#160;      using ::Eigen::internal::pload;</div>
<div class="line"><a name="l01344"></a><span class="lineno"> 1344</span>&#160;      using ::Eigen::internal::ploadt;</div>
<div class="line"><a name="l01345"></a><span class="lineno"> 1345</span>&#160;      using ::Eigen::internal::pstoret;</div>
<div class="line"><a name="l01346"></a><span class="lineno"> 1346</span>&#160; </div>
<div class="line"><a name="l01347"></a><span class="lineno"> 1347</span>&#160;      <span class="keyword">const</span> <span class="keywordtype">int</span> output_packet_size =</div>
<div class="line"><a name="l01348"></a><span class="lineno"> 1348</span>&#160;          internal::unpacket_traits&lt;PacketReturnType&gt;::size;</div>
<div class="line"><a name="l01349"></a><span class="lineno"> 1349</span>&#160; </div>
<div class="line"><a name="l01350"></a><span class="lineno"> 1350</span>&#160;      <span class="keywordtype">size_t</span> i = 0;</div>
<div class="line"><a name="l01351"></a><span class="lineno"> 1351</span>&#160;      <span class="keyword">const</span> <span class="keywordtype">size_t</span> num_packets = n / output_packet_size;</div>
<div class="line"><a name="l01352"></a><span class="lineno"> 1352</span>&#160;      <span class="keywordflow">for</span> (; i &lt; output_packet_size * num_packets; i += output_packet_size) {</div>
<div class="line"><a name="l01353"></a><span class="lineno"> 1353</span>&#160;        <span class="keyword">const</span> <span class="keyword">auto</span> src_val0 = pload&lt;PacketReturnType&gt;(src_buf0 + i);</div>
<div class="line"><a name="l01354"></a><span class="lineno"> 1354</span>&#160;        <span class="keyword">const</span> <span class="keyword">auto</span> src_val1 = pload&lt;PacketReturnType&gt;(src_buf1 + i);</div>
<div class="line"><a name="l01355"></a><span class="lineno"> 1355</span>&#160;        <span class="keyword">const</span> <span class="keyword">auto</span> src_val2 = pload&lt;PacketReturnType&gt;(src_buf2 + i);</div>
<div class="line"><a name="l01356"></a><span class="lineno"> 1356</span>&#160; </div>
<div class="line"><a name="l01357"></a><span class="lineno"> 1357</span>&#160;        <span class="keyword">const</span> <span class="keyword">auto</span> dst_val = ploadt&lt;PacketReturnType, Alignment&gt;(dst_buf + i);</div>
<div class="line"><a name="l01358"></a><span class="lineno"> 1358</span>&#160;        <span class="keyword">const</span> <span class="keyword">auto</span> sum =</div>
<div class="line"><a name="l01359"></a><span class="lineno"> 1359</span>&#160;            padd(padd(dst_val, src_val0), padd(src_val1, src_val2));</div>
<div class="line"><a name="l01360"></a><span class="lineno"> 1360</span>&#160; </div>
<div class="line"><a name="l01361"></a><span class="lineno"> 1361</span>&#160;        pstoret&lt;Scalar, PacketReturnType, Alignment&gt;(dst_buf + i, sum);</div>
<div class="line"><a name="l01362"></a><span class="lineno"> 1362</span>&#160;      }</div>
<div class="line"><a name="l01363"></a><span class="lineno"> 1363</span>&#160;      <span class="keywordflow">for</span> (; i &lt; n; ++i) {</div>
<div class="line"><a name="l01364"></a><span class="lineno"> 1364</span>&#160;        dst_buf[i] += src_buf0[i] + src_buf1[i] + src_buf2[i];</div>
<div class="line"><a name="l01365"></a><span class="lineno"> 1365</span>&#160;      }</div>
<div class="line"><a name="l01366"></a><span class="lineno"> 1366</span>&#160;    }</div>
<div class="line"><a name="l01367"></a><span class="lineno"> 1367</span>&#160; </div>
<div class="line"><a name="l01368"></a><span class="lineno"> 1368</span>&#160;    <span class="keyword">template</span> &lt;<span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l01369"></a><span class="lineno"> 1369</span>&#160;    <span class="keywordtype">void</span> eval(Barrier&amp; barrier, Index start_block_idx, Index end_block_idx) {</div>
<div class="line"><a name="l01370"></a><span class="lineno"> 1370</span>&#160;      <span class="keywordflow">while</span> (end_block_idx - start_block_idx &gt; 1) {</div>
<div class="line"><a name="l01371"></a><span class="lineno"> 1371</span>&#160;        <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> mid_block_idx = (start_block_idx + end_block_idx) / 2;</div>
<div class="line"><a name="l01372"></a><span class="lineno"> 1372</span>&#160;        evaluator-&gt;m_device.enqueueNoNotification(</div>
<div class="line"><a name="l01373"></a><span class="lineno"> 1373</span>&#160;            [<span class="keyword">this</span>, &amp;barrier, mid_block_idx, end_block_idx]() {</div>
<div class="line"><a name="l01374"></a><span class="lineno"> 1374</span>&#160;              eval&lt;Alignment&gt;(barrier, mid_block_idx, end_block_idx);</div>
<div class="line"><a name="l01375"></a><span class="lineno"> 1375</span>&#160;            });</div>
<div class="line"><a name="l01376"></a><span class="lineno"> 1376</span>&#160;        end_block_idx = mid_block_idx;</div>
<div class="line"><a name="l01377"></a><span class="lineno"> 1377</span>&#160;      }</div>
<div class="line"><a name="l01378"></a><span class="lineno"> 1378</span>&#160; </div>
<div class="line"><a name="l01379"></a><span class="lineno"> 1379</span>&#160;      <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> block_idx = start_block_idx;</div>
<div class="line"><a name="l01380"></a><span class="lineno"> 1380</span>&#160;      <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> block_start = block_idx * block_size;</div>
<div class="line"><a name="l01381"></a><span class="lineno"> 1381</span>&#160;      <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> block_end = block_start + actualBlockSize(block_idx);</div>
<div class="line"><a name="l01382"></a><span class="lineno"> 1382</span>&#160; </div>
<div class="line"><a name="l01383"></a><span class="lineno"> 1383</span>&#160;      processBlock&lt;Alignment&gt;(block_idx, block_start, block_end);</div>
<div class="line"><a name="l01384"></a><span class="lineno"> 1384</span>&#160;      barrier.Notify();</div>
<div class="line"><a name="l01385"></a><span class="lineno"> 1385</span>&#160;    }</div>
<div class="line"><a name="l01386"></a><span class="lineno"> 1386</span>&#160; </div>
<div class="line"><a name="l01387"></a><span class="lineno"> 1387</span>&#160;    <span class="keyword">template</span> &lt;<span class="keywordtype">int</span> Alignment&gt;</div>
<div class="line"><a name="l01388"></a><span class="lineno"> 1388</span>&#160;    <span class="keywordtype">void</span> evalAsync(Index start_block_idx, Index end_block_idx) {</div>
<div class="line"><a name="l01389"></a><span class="lineno"> 1389</span>&#160;      <span class="keywordflow">while</span> (end_block_idx - start_block_idx &gt; 1) {</div>
<div class="line"><a name="l01390"></a><span class="lineno"> 1390</span>&#160;        <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> mid_block_idx = (start_block_idx + end_block_idx) / 2;</div>
<div class="line"><a name="l01391"></a><span class="lineno"> 1391</span>&#160;        evaluator-&gt;m_device.enqueueNoNotification(</div>
<div class="line"><a name="l01392"></a><span class="lineno"> 1392</span>&#160;            [<span class="keyword">this</span>, mid_block_idx, end_block_idx]() {</div>
<div class="line"><a name="l01393"></a><span class="lineno"> 1393</span>&#160;              evalAsync&lt;Alignment&gt;(mid_block_idx, end_block_idx);</div>
<div class="line"><a name="l01394"></a><span class="lineno"> 1394</span>&#160;            });</div>
<div class="line"><a name="l01395"></a><span class="lineno"> 1395</span>&#160;        end_block_idx = mid_block_idx;</div>
<div class="line"><a name="l01396"></a><span class="lineno"> 1396</span>&#160;      }</div>
<div class="line"><a name="l01397"></a><span class="lineno"> 1397</span>&#160; </div>
<div class="line"><a name="l01398"></a><span class="lineno"> 1398</span>&#160;      <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> block_idx = start_block_idx;</div>
<div class="line"><a name="l01399"></a><span class="lineno"> 1399</span>&#160; </div>
<div class="line"><a name="l01400"></a><span class="lineno"> 1400</span>&#160;      <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> block_start = block_idx * block_size;</div>
<div class="line"><a name="l01401"></a><span class="lineno"> 1401</span>&#160;      <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> block_end = block_start + actualBlockSize(block_idx);</div>
<div class="line"><a name="l01402"></a><span class="lineno"> 1402</span>&#160; </div>
<div class="line"><a name="l01403"></a><span class="lineno"> 1403</span>&#160;      processBlock&lt;Alignment&gt;(block_idx, block_start, block_end);</div>
<div class="line"><a name="l01404"></a><span class="lineno"> 1404</span>&#160; </div>
<div class="line"><a name="l01405"></a><span class="lineno"> 1405</span>&#160;      <span class="keywordtype">int</span> v = num_pending_blocks.fetch_sub(1);</div>
<div class="line"><a name="l01406"></a><span class="lineno"> 1406</span>&#160;      eigen_assert(v &gt;= 1);</div>
<div class="line"><a name="l01407"></a><span class="lineno"> 1407</span>&#160; </div>
<div class="line"><a name="l01408"></a><span class="lineno"> 1408</span>&#160;      <span class="keywordflow">if</span> (v == 1) {</div>
<div class="line"><a name="l01409"></a><span class="lineno"> 1409</span>&#160;        <span class="comment">// Aggregate partial sums from l0 ranges.</span></div>
<div class="line"><a name="l01410"></a><span class="lineno"> 1410</span>&#160;        aggregateL0Blocks&lt;Alignment&gt;();</div>
<div class="line"><a name="l01411"></a><span class="lineno"> 1411</span>&#160; </div>
<div class="line"><a name="l01412"></a><span class="lineno"> 1412</span>&#160;        <span class="comment">// Apply output kernel.</span></div>
<div class="line"><a name="l01413"></a><span class="lineno"> 1413</span>&#160;        applyOutputKernel();</div>
<div class="line"><a name="l01414"></a><span class="lineno"> 1414</span>&#160; </div>
<div class="line"><a name="l01415"></a><span class="lineno"> 1415</span>&#160;        <span class="comment">// NOTE: If we call `done` callback before deleting this (context),</span></div>
<div class="line"><a name="l01416"></a><span class="lineno"> 1416</span>&#160;        <span class="comment">// it might deallocate Self* pointer captured by context, and we&#39;ll</span></div>
<div class="line"><a name="l01417"></a><span class="lineno"> 1417</span>&#160;        <span class="comment">// fail in destructor trying to deallocate temporary buffers.</span></div>
<div class="line"><a name="l01418"></a><span class="lineno"> 1418</span>&#160; </div>
<div class="line"><a name="l01419"></a><span class="lineno"> 1419</span>&#160;        <span class="comment">// Move done call back from context before it will be destructed.</span></div>
<div class="line"><a name="l01420"></a><span class="lineno"> 1420</span>&#160;        DoneCallback done_copy = std::move(done);</div>
<div class="line"><a name="l01421"></a><span class="lineno"> 1421</span>&#160; </div>
<div class="line"><a name="l01422"></a><span class="lineno"> 1422</span>&#160;        <span class="comment">// We are confident that we are the last one who touches context.</span></div>
<div class="line"><a name="l01423"></a><span class="lineno"> 1423</span>&#160;        <span class="keyword">delete</span> <span class="keyword">this</span>;</div>
<div class="line"><a name="l01424"></a><span class="lineno"> 1424</span>&#160; </div>
<div class="line"><a name="l01425"></a><span class="lineno"> 1425</span>&#160;        <span class="comment">// Now safely call the done callback.</span></div>
<div class="line"><a name="l01426"></a><span class="lineno"> 1426</span>&#160;        done_copy();</div>
<div class="line"><a name="l01427"></a><span class="lineno"> 1427</span>&#160;      }</div>
<div class="line"><a name="l01428"></a><span class="lineno"> 1428</span>&#160;    }</div>
<div class="line"><a name="l01429"></a><span class="lineno"> 1429</span>&#160; </div>
<div class="line"><a name="l01430"></a><span class="lineno"> 1430</span>&#160;    <span class="comment">// Cost model doesn&#39;t capture well the cost associated with constructing</span></div>
<div class="line"><a name="l01431"></a><span class="lineno"> 1431</span>&#160;    <span class="comment">// tensor contraction mappers and computing loop bounds in gemm_pack_lhs</span></div>
<div class="line"><a name="l01432"></a><span class="lineno"> 1432</span>&#160;    <span class="comment">// and gemm_pack_rhs, so we specify minimum desired block size.</span></div>
<div class="line"><a name="l01433"></a><span class="lineno"> 1433</span>&#160;    <span class="keyword">static</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> blockSize(Index k, <span class="keywordtype">int</span> num_threads) {</div>
<div class="line"><a name="l01434"></a><span class="lineno"> 1434</span>&#160;      <span class="keyword">const</span> <span class="keyword">auto</span> round_up = [=](<a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> index) -&gt; Index {</div>
<div class="line"><a name="l01435"></a><span class="lineno"> 1435</span>&#160;        <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> kmultiple = packet_size &lt;= 8 ? 8 : packet_size;</div>
<div class="line"><a name="l01436"></a><span class="lineno"> 1436</span>&#160;        <span class="keywordflow">return</span> divup&lt;Index&gt;(index, kmultiple) * kmultiple;</div>
<div class="line"><a name="l01437"></a><span class="lineno"> 1437</span>&#160;      };</div>
<div class="line"><a name="l01438"></a><span class="lineno"> 1438</span>&#160; </div>
<div class="line"><a name="l01439"></a><span class="lineno"> 1439</span>&#160;      <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> target_block_size = round_up(divup&lt;Index&gt;(k, num_threads));</div>
<div class="line"><a name="l01440"></a><span class="lineno"> 1440</span>&#160;      <span class="keyword">const</span> <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> desired_min_block_size = 12 * packet_size;</div>
<div class="line"><a name="l01441"></a><span class="lineno"> 1441</span>&#160; </div>
<div class="line"><a name="l01442"></a><span class="lineno"> 1442</span>&#160;      <span class="keywordflow">return</span> numext::mini&lt;Index&gt;(</div>
<div class="line"><a name="l01443"></a><span class="lineno"> 1443</span>&#160;          k, numext::maxi&lt;Index&gt;(desired_min_block_size, target_block_size));</div>
<div class="line"><a name="l01444"></a><span class="lineno"> 1444</span>&#160;    }</div>
<div class="line"><a name="l01445"></a><span class="lineno"> 1445</span>&#160; </div>
<div class="line"><a name="l01446"></a><span class="lineno"> 1446</span>&#160;    EvalShardedByInnerDimContext(<span class="keyword">const</span> EvalShardedByInnerDimContext&amp;) = <span class="keyword">delete</span>;</div>
<div class="line"><a name="l01447"></a><span class="lineno"> 1447</span>&#160;    <span class="keywordtype">void</span> operator=(<span class="keyword">const</span> EvalShardedByInnerDimContext&amp;) = <span class="keyword">delete</span>;</div>
<div class="line"><a name="l01448"></a><span class="lineno"> 1448</span>&#160;  };</div>
<div class="line"><a name="l01449"></a><span class="lineno"> 1449</span>&#160; </div>
<div class="line"><a name="l01450"></a><span class="lineno"> 1450</span>&#160;  <span class="comment">// ------------------------------------------------------------------------ //</span></div>
<div class="line"><a name="l01451"></a><span class="lineno"> 1451</span>&#160; </div>
<div class="line"><a name="l01452"></a><span class="lineno"> 1452</span>&#160;  <span class="comment">// Below are the function used by evalProductImpl heuristics, trying to select</span></div>
<div class="line"><a name="l01453"></a><span class="lineno"> 1453</span>&#160;  <span class="comment">// optimcal parameters for parallelization algorithm.</span></div>
<div class="line"><a name="l01454"></a><span class="lineno"> 1454</span>&#160; </div>
<div class="line"><a name="l01455"></a><span class="lineno"> 1455</span>&#160;  <span class="comment">// Decide whether we want to shard m x n contraction by columns or by rows.</span></div>
<div class="line"><a name="l01456"></a><span class="lineno"> 1456</span>&#160;  <span class="keyword">static</span> <span class="keywordtype">bool</span> shardByCol(Index m, Index n, Index num_threads) {</div>
<div class="line"><a name="l01457"></a><span class="lineno"> 1457</span>&#160;    <span class="comment">// Note: we are comparing both n and m against Traits::nr, it is not</span></div>
<div class="line"><a name="l01458"></a><span class="lineno"> 1458</span>&#160;    <span class="comment">// a mistake. We are trying to figure out how both n and m will fit into</span></div>
<div class="line"><a name="l01459"></a><span class="lineno"> 1459</span>&#160;    <span class="comment">// the main sharding dimension.</span></div>
<div class="line"><a name="l01460"></a><span class="lineno"> 1460</span>&#160; </div>
<div class="line"><a name="l01461"></a><span class="lineno"> 1461</span>&#160;    <span class="comment">// Sharding by column is the default</span></div>
<div class="line"><a name="l01462"></a><span class="lineno"> 1462</span>&#160;    <span class="comment">// ... unless there is enough data for vectorization over rows</span></div>
<div class="line"><a name="l01463"></a><span class="lineno"> 1463</span>&#160;    <span class="keywordflow">if</span> (m / num_threads &gt;= Traits::nr &amp;&amp;</div>
<div class="line"><a name="l01464"></a><span class="lineno"> 1464</span>&#160;        <span class="comment">// and not enough data for vectorization over columns</span></div>
<div class="line"><a name="l01465"></a><span class="lineno"> 1465</span>&#160;        (n / num_threads &lt; Traits::nr ||</div>
<div class="line"><a name="l01466"></a><span class="lineno"> 1466</span>&#160;         <span class="comment">// ... or barely enough data for vectorization over columns,</span></div>
<div class="line"><a name="l01467"></a><span class="lineno"> 1467</span>&#160;         <span class="comment">// but it is not evenly dividable across threads</span></div>
<div class="line"><a name="l01468"></a><span class="lineno"> 1468</span>&#160;         (n / num_threads &lt; 4 * Traits::nr &amp;&amp;</div>
<div class="line"><a name="l01469"></a><span class="lineno"> 1469</span>&#160;          (n % (num_threads * Traits::nr)) != 0 &amp;&amp;</div>
<div class="line"><a name="l01470"></a><span class="lineno"> 1470</span>&#160;          <span class="comment">// ... and it is evenly dividable across threads for rows</span></div>
<div class="line"><a name="l01471"></a><span class="lineno"> 1471</span>&#160;          ((m % (num_threads * Traits::nr)) == 0 ||</div>
<div class="line"><a name="l01472"></a><span class="lineno"> 1472</span>&#160;           <span class="comment">// .. or it is not evenly dividable for both dimensions but</span></div>
<div class="line"><a name="l01473"></a><span class="lineno"> 1473</span>&#160;           <span class="comment">// there is much more data over rows so that corner effects are</span></div>
<div class="line"><a name="l01474"></a><span class="lineno"> 1474</span>&#160;           <span class="comment">// mitigated.</span></div>
<div class="line"><a name="l01475"></a><span class="lineno"> 1475</span>&#160;           (m / n &gt;= 6)))))</div>
<div class="line"><a name="l01476"></a><span class="lineno"> 1476</span>&#160;      <span class="keywordflow">return</span> <span class="keyword">false</span>;</div>
<div class="line"><a name="l01477"></a><span class="lineno"> 1477</span>&#160;    <span class="comment">// Wait, or if matrices are just substantially prolonged over the other</span></div>
<div class="line"><a name="l01478"></a><span class="lineno"> 1478</span>&#160;    <span class="comment">// dimension.</span></div>
<div class="line"><a name="l01479"></a><span class="lineno"> 1479</span>&#160;    <span class="keywordflow">if</span> (n / num_threads &lt; 16 * Traits::nr &amp;&amp; m &gt; n * 32) <span class="keywordflow">return</span> <span class="keyword">false</span>;</div>
<div class="line"><a name="l01480"></a><span class="lineno"> 1480</span>&#160;    <span class="keywordflow">return</span> <span class="keyword">true</span>;</div>
<div class="line"><a name="l01481"></a><span class="lineno"> 1481</span>&#160;  }</div>
<div class="line"><a name="l01482"></a><span class="lineno"> 1482</span>&#160; </div>
<div class="line"><a name="l01483"></a><span class="lineno"> 1483</span>&#160;  <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> coarsenM(Index m, Index n, Index bm, Index bn, Index bk, Index gn,</div>
<div class="line"><a name="l01484"></a><span class="lineno"> 1484</span>&#160;                 <span class="keywordtype">int</span> num_threads, <span class="keywordtype">bool</span> shard_by_col)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01485"></a><span class="lineno"> 1485</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> gm = 1;</div>
<div class="line"><a name="l01486"></a><span class="lineno"> 1486</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> gm1 = 1;</div>
<div class="line"><a name="l01487"></a><span class="lineno"> 1487</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nm0 = divup(m, bm);</div>
<div class="line"><a name="l01488"></a><span class="lineno"> 1488</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nm1 = nm0;</div>
<div class="line"><a name="l01489"></a><span class="lineno"> 1489</span>&#160;    <span class="keywordflow">for</span> (;;) {</div>
<div class="line"><a name="l01490"></a><span class="lineno"> 1490</span>&#160;      <span class="comment">// Find the next candidate for m grain size. It needs to result in</span></div>
<div class="line"><a name="l01491"></a><span class="lineno"> 1491</span>&#160;      <span class="comment">// different number of blocks. E.g. if we have 10 kernels, we want to try</span></div>
<div class="line"><a name="l01492"></a><span class="lineno"> 1492</span>&#160;      <span class="comment">// 5 and 10, but not 6, 7, 8 and 9.</span></div>
<div class="line"><a name="l01493"></a><span class="lineno"> 1493</span>&#160;      <span class="keywordflow">while</span> (gm1 &lt;= nm0 &amp;&amp; nm1 == divup(nm0, gm1)) gm1++;</div>
<div class="line"><a name="l01494"></a><span class="lineno"> 1494</span>&#160;      <span class="keywordflow">if</span> (gm1 &gt; nm0) <span class="keywordflow">break</span>;</div>
<div class="line"><a name="l01495"></a><span class="lineno"> 1495</span>&#160;      <span class="comment">// Check the candidate.</span></div>
<div class="line"><a name="l01496"></a><span class="lineno"> 1496</span>&#160;      <span class="keywordtype">int</span> res = checkGrain(m, n, bm, bn, bk, gm1, gn, gm, gn, num_threads,</div>
<div class="line"><a name="l01497"></a><span class="lineno"> 1497</span>&#160;                           shard_by_col);</div>
<div class="line"><a name="l01498"></a><span class="lineno"> 1498</span>&#160;      <span class="keywordflow">if</span> (res &lt; 0) <span class="keywordflow">break</span>;</div>
<div class="line"><a name="l01499"></a><span class="lineno"> 1499</span>&#160;      nm1 = divup(nm0, gm1);</div>
<div class="line"><a name="l01500"></a><span class="lineno"> 1500</span>&#160;      <span class="keywordflow">if</span> (res == 0) <span class="keywordflow">continue</span>;</div>
<div class="line"><a name="l01501"></a><span class="lineno"> 1501</span>&#160;      <span class="comment">// Commit new grain size.</span></div>
<div class="line"><a name="l01502"></a><span class="lineno"> 1502</span>&#160;      gm = gm1;</div>
<div class="line"><a name="l01503"></a><span class="lineno"> 1503</span>&#160;    }</div>
<div class="line"><a name="l01504"></a><span class="lineno"> 1504</span>&#160;    <span class="keywordflow">return</span> gm;</div>
<div class="line"><a name="l01505"></a><span class="lineno"> 1505</span>&#160;  }</div>
<div class="line"><a name="l01506"></a><span class="lineno"> 1506</span>&#160; </div>
<div class="line"><a name="l01507"></a><span class="lineno"> 1507</span>&#160;  <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> coarsenN(Index m, Index n, Index bm, Index bn, Index bk, Index gm,</div>
<div class="line"><a name="l01508"></a><span class="lineno"> 1508</span>&#160;                 <span class="keywordtype">int</span> num_threads, <span class="keywordtype">bool</span> shard_by_col)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01509"></a><span class="lineno"> 1509</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> gn = 1;</div>
<div class="line"><a name="l01510"></a><span class="lineno"> 1510</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> gn1 = 1;</div>
<div class="line"><a name="l01511"></a><span class="lineno"> 1511</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nn0 = divup(n, bn);</div>
<div class="line"><a name="l01512"></a><span class="lineno"> 1512</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nn1 = nn0;</div>
<div class="line"><a name="l01513"></a><span class="lineno"> 1513</span>&#160;    <span class="keywordflow">for</span> (;;) {</div>
<div class="line"><a name="l01514"></a><span class="lineno"> 1514</span>&#160;      <span class="keywordflow">while</span> (gn1 &lt;= nn0 &amp;&amp; nn1 == divup(nn0, gn1)) gn1++;</div>
<div class="line"><a name="l01515"></a><span class="lineno"> 1515</span>&#160;      <span class="keywordflow">if</span> (gn1 &gt; nn0) <span class="keywordflow">break</span>;</div>
<div class="line"><a name="l01516"></a><span class="lineno"> 1516</span>&#160;      <span class="keywordtype">int</span> res = checkGrain(m, n, bm, bn, bk, gm, gn1, gm, gn, num_threads,</div>
<div class="line"><a name="l01517"></a><span class="lineno"> 1517</span>&#160;                           shard_by_col);</div>
<div class="line"><a name="l01518"></a><span class="lineno"> 1518</span>&#160;      <span class="keywordflow">if</span> (res &lt; 0) <span class="keywordflow">break</span>;</div>
<div class="line"><a name="l01519"></a><span class="lineno"> 1519</span>&#160;      nn1 = divup(nn0, gn1);</div>
<div class="line"><a name="l01520"></a><span class="lineno"> 1520</span>&#160;      <span class="keywordflow">if</span> (res == 0) <span class="keywordflow">continue</span>;</div>
<div class="line"><a name="l01521"></a><span class="lineno"> 1521</span>&#160;      gn = gn1;</div>
<div class="line"><a name="l01522"></a><span class="lineno"> 1522</span>&#160;    }</div>
<div class="line"><a name="l01523"></a><span class="lineno"> 1523</span>&#160;    <span class="keywordflow">return</span> gn;</div>
<div class="line"><a name="l01524"></a><span class="lineno"> 1524</span>&#160;  }</div>
<div class="line"><a name="l01525"></a><span class="lineno"> 1525</span>&#160; </div>
<div class="line"><a name="l01526"></a><span class="lineno"> 1526</span>&#160;  <span class="comment">// checkGrain checks whether grain (gm, gn) is suitable and is better than</span></div>
<div class="line"><a name="l01527"></a><span class="lineno"> 1527</span>&#160;  <span class="comment">// (oldgm, oldgn).</span></div>
<div class="line"><a name="l01528"></a><span class="lineno"> 1528</span>&#160;  <span class="keywordtype">int</span> checkGrain(Index m, Index n, Index bm, Index bn, Index bk, Index gm,</div>
<div class="line"><a name="l01529"></a><span class="lineno"> 1529</span>&#160;                 Index gn, Index oldgm, Index oldgn, <span class="keywordtype">int</span> num_threads,</div>
<div class="line"><a name="l01530"></a><span class="lineno"> 1530</span>&#160;                 <span class="keywordtype">bool</span> shard_by_col)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01531"></a><span class="lineno"> 1531</span>&#160;    <span class="keyword">const</span> TensorOpCost cost =</div>
<div class="line"><a name="l01532"></a><span class="lineno"> 1532</span>&#160;        contractionCost(bm * gm, bn * gn, bm, bn, bk, shard_by_col, <span class="keyword">true</span>);</div>
<div class="line"><a name="l01533"></a><span class="lineno"> 1533</span>&#160;    <span class="keywordtype">double</span> taskSize = TensorCostModel&lt;ThreadPoolDevice&gt;::taskSize(</div>
<div class="line"><a name="l01534"></a><span class="lineno"> 1534</span>&#160;        <span class="keyword">static_cast&lt;</span><span class="keywordtype">double</span><span class="keyword">&gt;</span>(bm) * gm * bn * gn, cost);</div>
<div class="line"><a name="l01535"></a><span class="lineno"> 1535</span>&#160;    <span class="comment">// If the task is too small, then we agree on it regardless of anything</span></div>
<div class="line"><a name="l01536"></a><span class="lineno"> 1536</span>&#160;    <span class="comment">// else. Otherwise synchronization overheads will dominate.</span></div>
<div class="line"><a name="l01537"></a><span class="lineno"> 1537</span>&#160;    <span class="keywordflow">if</span> (taskSize &lt; 1) <span class="keywordflow">return</span> 1;</div>
<div class="line"><a name="l01538"></a><span class="lineno"> 1538</span>&#160;    <span class="comment">// If it is too large, then we reject it and all larger tasks.</span></div>
<div class="line"><a name="l01539"></a><span class="lineno"> 1539</span>&#160;    <span class="keywordflow">if</span> (taskSize &gt; 2) <span class="keywordflow">return</span> -1;</div>
<div class="line"><a name="l01540"></a><span class="lineno"> 1540</span>&#160;    <span class="comment">// Now we are in presumably good task size range.</span></div>
<div class="line"><a name="l01541"></a><span class="lineno"> 1541</span>&#160;    <span class="comment">// The main deciding factor here is parallelism. Consider that we have 12</span></div>
<div class="line"><a name="l01542"></a><span class="lineno"> 1542</span>&#160;    <span class="comment">// kernels and 4 threads. Grains of 2, 3 and 4 all yield good task sizes.</span></div>
<div class="line"><a name="l01543"></a><span class="lineno"> 1543</span>&#160;    <span class="comment">// But 2/4 yield 6/3 tasks, which gives us parallelism of 0.75 (at most 3/4</span></div>
<div class="line"><a name="l01544"></a><span class="lineno"> 1544</span>&#160;    <span class="comment">// of cores will be busy). While grain size 3 gives us 4 tasks, which gives</span></div>
<div class="line"><a name="l01545"></a><span class="lineno"> 1545</span>&#160;    <span class="comment">// us parallelism of 1 (we can load all cores).</span></div>
<div class="line"><a name="l01546"></a><span class="lineno"> 1546</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nm0 = divup(m, bm);</div>
<div class="line"><a name="l01547"></a><span class="lineno"> 1547</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> nn0 = divup(n, bn);</div>
<div class="line"><a name="l01548"></a><span class="lineno"> 1548</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> new_tasks = divup(nm0, gm) * divup(nn0, gn);</div>
<div class="line"><a name="l01549"></a><span class="lineno"> 1549</span>&#160;    <span class="keywordtype">double</span> new_parallelism = <span class="keyword">static_cast&lt;</span><span class="keywordtype">double</span><span class="keyword">&gt;</span>(new_tasks) /</div>
<div class="line"><a name="l01550"></a><span class="lineno"> 1550</span>&#160;                             (divup&lt;int&gt;(new_tasks, num_threads) * num_threads);</div>
<div class="line"><a name="l01551"></a><span class="lineno"> 1551</span>&#160;    <a class="codeRef" href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Index</a> old_tasks = divup(nm0, oldgm) * divup(nn0, oldgn);</div>
<div class="line"><a name="l01552"></a><span class="lineno"> 1552</span>&#160;    <span class="keywordtype">double</span> old_parallelism = <span class="keyword">static_cast&lt;</span><span class="keywordtype">double</span><span class="keyword">&gt;</span>(old_tasks) /</div>
<div class="line"><a name="l01553"></a><span class="lineno"> 1553</span>&#160;                             (divup&lt;int&gt;(old_tasks, num_threads) * num_threads);</div>
<div class="line"><a name="l01554"></a><span class="lineno"> 1554</span>&#160;    <span class="keywordflow">if</span> (new_parallelism &gt; old_parallelism || new_parallelism == 1) <span class="keywordflow">return</span> 1;</div>
<div class="line"><a name="l01555"></a><span class="lineno"> 1555</span>&#160;    <span class="keywordflow">return</span> 0;</div>
<div class="line"><a name="l01556"></a><span class="lineno"> 1556</span>&#160;  }</div>
<div class="line"><a name="l01557"></a><span class="lineno"> 1557</span>&#160; </div>
<div class="line"><a name="l01558"></a><span class="lineno"> 1558</span>&#160;  TensorOpCost contractionCost(Index m, Index n, Index bm, Index bn, Index bk,</div>
<div class="line"><a name="l01559"></a><span class="lineno"> 1559</span>&#160;                               <span class="keywordtype">bool</span> shard_by_col, <span class="keywordtype">bool</span> prepacked)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01560"></a><span class="lineno"> 1560</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">int</span> packed_size = std::min&lt;int&gt;(PacketType&lt;LhsScalar, Device&gt;::size,</div>
<div class="line"><a name="l01561"></a><span class="lineno"> 1561</span>&#160;                                          PacketType&lt;RhsScalar, Device&gt;::size);</div>
<div class="line"><a name="l01562"></a><span class="lineno"> 1562</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">int</span> output_packet_size = internal::unpacket_traits&lt;PacketReturnType&gt;::size;</div>
<div class="line"><a name="l01563"></a><span class="lineno"> 1563</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">double</span> kd = <span class="keyword">static_cast&lt;</span><span class="keywordtype">double</span><span class="keyword">&gt;</span>(bk);</div>
<div class="line"><a name="l01564"></a><span class="lineno"> 1564</span>&#160;    <span class="keywordtype">double</span> compute_bandwidth = computeBandwidth(<span class="keyword">false</span>, bm, bn, bk);</div>
<div class="line"><a name="l01565"></a><span class="lineno"> 1565</span>&#160;    <span class="comment">// Computations.</span></div>
<div class="line"><a name="l01566"></a><span class="lineno"> 1566</span>&#160;    TensorOpCost cost = TensorOpCost(0, 0, kd * compute_bandwidth, <span class="keyword">true</span>, packed_size);</div>
<div class="line"><a name="l01567"></a><span class="lineno"> 1567</span>&#160;    <span class="comment">// Output stores.</span></div>
<div class="line"><a name="l01568"></a><span class="lineno"> 1568</span>&#160;    cost += TensorOpCost(0, <span class="keyword">sizeof</span>(CoeffReturnType), 0, <span class="keyword">true</span>, output_packet_size);</div>
<div class="line"><a name="l01569"></a><span class="lineno"> 1569</span>&#160;    <span class="keywordflow">if</span> (prepacked) {</div>
<div class="line"><a name="l01570"></a><span class="lineno"> 1570</span>&#160;      <span class="comment">// Packing and kernels are executed in different tasks. When we calculate</span></div>
<div class="line"><a name="l01571"></a><span class="lineno"> 1571</span>&#160;      <span class="comment">// task grain size we look only at kernel cost assuming that kernel</span></div>
<div class="line"><a name="l01572"></a><span class="lineno"> 1572</span>&#160;      <span class="comment">// is more expensive than packing.</span></div>
<div class="line"><a name="l01573"></a><span class="lineno"> 1573</span>&#160;      <span class="keywordflow">return</span> cost;</div>
<div class="line"><a name="l01574"></a><span class="lineno"> 1574</span>&#160;    }</div>
<div class="line"><a name="l01575"></a><span class="lineno"> 1575</span>&#160;    <span class="comment">// Lhs/rhs loads + computations.</span></div>
<div class="line"><a name="l01576"></a><span class="lineno"> 1576</span>&#160;    TensorOpCost lhsCost = this-&gt;m_leftImpl.costPerCoeff(<span class="keyword">true</span>) * (kd / n);</div>
<div class="line"><a name="l01577"></a><span class="lineno"> 1577</span>&#160;    TensorOpCost rhsCost = this-&gt;m_rightImpl.costPerCoeff(<span class="keyword">true</span>) * (kd / m);</div>
<div class="line"><a name="l01578"></a><span class="lineno"> 1578</span>&#160;    <span class="comment">// Lhs packing memory cost does not contribute considerably to overall</span></div>
<div class="line"><a name="l01579"></a><span class="lineno"> 1579</span>&#160;    <span class="comment">// execution time because lhs is prefetched early and accessed sequentially.</span></div>
<div class="line"><a name="l01580"></a><span class="lineno"> 1580</span>&#160;    <span class="keywordflow">if</span> (shard_by_col)</div>
<div class="line"><a name="l01581"></a><span class="lineno"> 1581</span>&#160;      lhsCost.dropMemoryCost();</div>
<div class="line"><a name="l01582"></a><span class="lineno"> 1582</span>&#160;    <span class="keywordflow">else</span></div>
<div class="line"><a name="l01583"></a><span class="lineno"> 1583</span>&#160;      rhsCost.dropMemoryCost();</div>
<div class="line"><a name="l01584"></a><span class="lineno"> 1584</span>&#160;    <span class="keywordflow">return</span> cost + lhsCost + rhsCost;</div>
<div class="line"><a name="l01585"></a><span class="lineno"> 1585</span>&#160;  }</div>
<div class="line"><a name="l01586"></a><span class="lineno"> 1586</span>&#160; </div>
<div class="line"><a name="l01587"></a><span class="lineno"> 1587</span>&#160;  <span class="comment">// Decide whether we want to shard m x k x n contraction over the inner</span></div>
<div class="line"><a name="l01588"></a><span class="lineno"> 1588</span>&#160;  <span class="comment">// (contraction) dimension (k).</span></div>
<div class="line"><a name="l01589"></a><span class="lineno"> 1589</span>&#160;  <span class="keyword">static</span> <span class="keywordtype">bool</span> shardByInnerDim(Index m, Index n, Index k, <span class="keywordtype">int</span> num_threads,</div>
<div class="line"><a name="l01590"></a><span class="lineno"> 1590</span>&#160;                              <span class="keywordtype">int</span> num_threads_by_k) {</div>
<div class="line"><a name="l01591"></a><span class="lineno"> 1591</span>&#160;    std::ptrdiff_t bufsize = m * n * <span class="keyword">sizeof</span>(Scalar);</div>
<div class="line"><a name="l01592"></a><span class="lineno"> 1592</span>&#160;    <span class="keywordtype">bool</span> shard_by_k = <span class="keyword">false</span>;</div>
<div class="line"><a name="l01593"></a><span class="lineno"> 1593</span>&#160;    <span class="keywordflow">if</span> (n == 1 ||                <span class="comment">// If mat*vec or...</span></div>
<div class="line"><a name="l01594"></a><span class="lineno"> 1594</span>&#160;        num_threads_by_k &lt; 2 ||  <span class="comment">// running single threaded or...</span></div>
<div class="line"><a name="l01595"></a><span class="lineno"> 1595</span>&#160;        num_threads_by_k &lt;</div>
<div class="line"><a name="l01596"></a><span class="lineno"> 1596</span>&#160;            num_threads ||  <span class="comment">// sharding by k gives less parallelism or...</span></div>
<div class="line"><a name="l01597"></a><span class="lineno"> 1597</span>&#160;        bufsize &gt; l3CacheSize() / num_threads_by_k ||  <span class="comment">// need more buffer space</span></div>
<div class="line"><a name="l01598"></a><span class="lineno"> 1598</span>&#160;        <span class="comment">// than L3 cache or...</span></div>
<div class="line"><a name="l01599"></a><span class="lineno"> 1599</span>&#160;        k / num_threads_by_k &lt; 2 * Traits::nr) {  <span class="comment">// k per thread is tiny.</span></div>
<div class="line"><a name="l01600"></a><span class="lineno"> 1600</span>&#160;      shard_by_k = <span class="keyword">false</span>;</div>
<div class="line"><a name="l01601"></a><span class="lineno"> 1601</span>&#160;    } <span class="keywordflow">else</span> <span class="keywordflow">if</span> (numext::maxi(m, n) / num_threads &lt;</div>
<div class="line"><a name="l01602"></a><span class="lineno"> 1602</span>&#160;                   Traits::nr ||  <span class="comment">// both other dimensions are tiny or...</span></div>
<div class="line"><a name="l01603"></a><span class="lineno"> 1603</span>&#160;               <span class="comment">// k per thread is not small and...</span></div>
<div class="line"><a name="l01604"></a><span class="lineno"> 1604</span>&#160;               (k / num_threads_by_k &gt; 8 * Traits::nr &amp;&amp;</div>
<div class="line"><a name="l01605"></a><span class="lineno"> 1605</span>&#160;                <span class="comment">// one of the outer dimensions is tiny or sharding by k offers</span></div>
<div class="line"><a name="l01606"></a><span class="lineno"> 1606</span>&#160;                <span class="comment">// more parallelism.</span></div>
<div class="line"><a name="l01607"></a><span class="lineno"> 1607</span>&#160;                (numext::mini(m, n) &lt; 2 * Traits::nr ||</div>
<div class="line"><a name="l01608"></a><span class="lineno"> 1608</span>&#160;                 num_threads_by_k &gt; num_threads))) {</div>
<div class="line"><a name="l01609"></a><span class="lineno"> 1609</span>&#160;      shard_by_k = <span class="keyword">true</span>;</div>
<div class="line"><a name="l01610"></a><span class="lineno"> 1610</span>&#160;    }</div>
<div class="line"><a name="l01611"></a><span class="lineno"> 1611</span>&#160;    <span class="keywordflow">return</span> shard_by_k;</div>
<div class="line"><a name="l01612"></a><span class="lineno"> 1612</span>&#160;  }</div>
<div class="line"><a name="l01613"></a><span class="lineno"> 1613</span>&#160; </div>
<div class="line"><a name="l01614"></a><span class="lineno"> 1614</span>&#160;  TensorOpCost contractionCostPerInnerDim(Index m, Index n, Index k)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01615"></a><span class="lineno"> 1615</span>&#160;    <span class="comment">// Compute cost.</span></div>
<div class="line"><a name="l01616"></a><span class="lineno"> 1616</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">int</span> output_packet_size = internal::unpacket_traits&lt;PacketReturnType&gt;::size;</div>
<div class="line"><a name="l01617"></a><span class="lineno"> 1617</span>&#160;    TensorOpCost cost(0, 0, (computeBandwidth(<span class="keyword">true</span>, m, n, k) * m) * n, <span class="keyword">true</span>, output_packet_size);</div>
<div class="line"><a name="l01618"></a><span class="lineno"> 1618</span>&#160;    <span class="comment">// Output stores.</span></div>
<div class="line"><a name="l01619"></a><span class="lineno"> 1619</span>&#160;    cost += TensorOpCost(0, <span class="keyword">sizeof</span>(CoeffReturnType), 0, <span class="keyword">true</span>, output_packet_size);</div>
<div class="line"><a name="l01620"></a><span class="lineno"> 1620</span>&#160;    TensorOpCost lhsCost = this-&gt;m_leftImpl.costPerCoeff(<span class="keyword">true</span>) * m;</div>
<div class="line"><a name="l01621"></a><span class="lineno"> 1621</span>&#160;    TensorOpCost rhsCost = this-&gt;m_rightImpl.costPerCoeff(<span class="keyword">true</span>) * n;</div>
<div class="line"><a name="l01622"></a><span class="lineno"> 1622</span>&#160;    <span class="comment">// Since the inner gemm kernel is always sharded by column, the lhs</span></div>
<div class="line"><a name="l01623"></a><span class="lineno"> 1623</span>&#160;    <span class="comment">// load cost is negligible.</span></div>
<div class="line"><a name="l01624"></a><span class="lineno"> 1624</span>&#160;    lhsCost.dropMemoryCost();</div>
<div class="line"><a name="l01625"></a><span class="lineno"> 1625</span>&#160;    <span class="keywordflow">return</span> cost + lhsCost + rhsCost;</div>
<div class="line"><a name="l01626"></a><span class="lineno"> 1626</span>&#160;  }</div>
<div class="line"><a name="l01627"></a><span class="lineno"> 1627</span>&#160; </div>
<div class="line"><a name="l01628"></a><span class="lineno"> 1628</span>&#160;  <span class="keywordtype">int</span> numThreadsInnerDim(Index m, Index n, Index k)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01629"></a><span class="lineno"> 1629</span>&#160;    <span class="keyword">const</span> <span class="keywordtype">int</span> output_packet_size = internal::unpacket_traits&lt;PacketReturnType&gt;::size;</div>
<div class="line"><a name="l01630"></a><span class="lineno"> 1630</span>&#160;    TensorOpCost cost = contractionCostPerInnerDim(m, n, k);</div>
<div class="line"><a name="l01631"></a><span class="lineno"> 1631</span>&#160;    <span class="keywordtype">double</span> total_parallel_cost =</div>
<div class="line"><a name="l01632"></a><span class="lineno"> 1632</span>&#160;        TensorCostModel&lt;ThreadPoolDevice&gt;::totalCost(k, cost);</div>
<div class="line"><a name="l01633"></a><span class="lineno"> 1633</span>&#160;    <span class="comment">// Cost of reduction step accumulating the m*n per-thread buffers into the</span></div>
<div class="line"><a name="l01634"></a><span class="lineno"> 1634</span>&#160;    <span class="comment">// result.</span></div>
<div class="line"><a name="l01635"></a><span class="lineno"> 1635</span>&#160;    <span class="keywordtype">double</span> reduction_cost = TensorCostModel&lt;ThreadPoolDevice&gt;::totalCost(</div>
<div class="line"><a name="l01636"></a><span class="lineno"> 1636</span>&#160;        m * n, TensorOpCost(2, 1, 1, <span class="keyword">true</span>, output_packet_size));</div>
<div class="line"><a name="l01637"></a><span class="lineno"> 1637</span>&#160;    <span class="keywordtype">int</span> num_threads = 1;</div>
<div class="line"><a name="l01638"></a><span class="lineno"> 1638</span>&#160;    <span class="keywordtype">double</span> min_cost = total_parallel_cost;</div>
<div class="line"><a name="l01639"></a><span class="lineno"> 1639</span>&#160;    <span class="keywordtype">double</span> kPerThreadOverHead = 3000;</div>
<div class="line"><a name="l01640"></a><span class="lineno"> 1640</span>&#160;    <span class="keywordtype">double</span> kFixedOverHead = 100000;</div>
<div class="line"><a name="l01641"></a><span class="lineno"> 1641</span>&#160;    <span class="keywordflow">for</span> (<span class="keywordtype">int</span> nt = 2; nt &lt;= this-&gt;m_device.numThreads(); nt += 2) {</div>
<div class="line"><a name="l01642"></a><span class="lineno"> 1642</span>&#160;      <span class="keywordtype">double</span> sequential_cost =</div>
<div class="line"><a name="l01643"></a><span class="lineno"> 1643</span>&#160;          kFixedOverHead + nt * (reduction_cost + kPerThreadOverHead);</div>
<div class="line"><a name="l01644"></a><span class="lineno"> 1644</span>&#160;      <span class="keywordtype">double</span> parallel_cost = total_parallel_cost / nt + sequential_cost;</div>
<div class="line"><a name="l01645"></a><span class="lineno"> 1645</span>&#160;      <span class="keywordflow">if</span> (parallel_cost &lt; min_cost) {</div>
<div class="line"><a name="l01646"></a><span class="lineno"> 1646</span>&#160;        num_threads = nt;</div>
<div class="line"><a name="l01647"></a><span class="lineno"> 1647</span>&#160;        min_cost = parallel_cost;</div>
<div class="line"><a name="l01648"></a><span class="lineno"> 1648</span>&#160;      }</div>
<div class="line"><a name="l01649"></a><span class="lineno"> 1649</span>&#160;    }</div>
<div class="line"><a name="l01650"></a><span class="lineno"> 1650</span>&#160;    <span class="keywordflow">return</span> num_threads;</div>
<div class="line"><a name="l01651"></a><span class="lineno"> 1651</span>&#160;  }</div>
<div class="line"><a name="l01652"></a><span class="lineno"> 1652</span>&#160; </div>
<div class="line"><a name="l01653"></a><span class="lineno"> 1653</span>&#160;  <span class="keywordtype">double</span> computeBandwidth(<span class="keywordtype">bool</span> shard_by_col, Index bm, Index bn,</div>
<div class="line"><a name="l01654"></a><span class="lineno"> 1654</span>&#160;                          Index bk)<span class="keyword"> const </span>{</div>
<div class="line"><a name="l01655"></a><span class="lineno"> 1655</span>&#160;    <span class="comment">// Peak VFMA bandwidth is 0.5. However if we have not enough data for</span></div>
<div class="line"><a name="l01656"></a><span class="lineno"> 1656</span>&#160;    <span class="comment">// vectorization bandwidth drops. The 4.0 and 2.0 bandwidth is determined</span></div>
<div class="line"><a name="l01657"></a><span class="lineno"> 1657</span>&#160;    <span class="comment">// experimentally.</span></div>
<div class="line"><a name="l01658"></a><span class="lineno"> 1658</span>&#160;    <span class="keywordtype">double</span> computeBandwidth =</div>
<div class="line"><a name="l01659"></a><span class="lineno"> 1659</span>&#160;        bk == 1 ? 4.0</div>
<div class="line"><a name="l01660"></a><span class="lineno"> 1660</span>&#160;                : (shard_by_col ? bn : bm) &lt; Traits::nr ||</div>
<div class="line"><a name="l01661"></a><span class="lineno"> 1661</span>&#160;                          (shard_by_col ? bm : bn) &lt; Traits::mr</div>
<div class="line"><a name="l01662"></a><span class="lineno"> 1662</span>&#160;                      ? 2.0</div>
<div class="line"><a name="l01663"></a><span class="lineno"> 1663</span>&#160;                      : 0.5;</div>
<div class="line"><a name="l01664"></a><span class="lineno"> 1664</span>&#160;<span class="preprocessor">#ifndef EIGEN_VECTORIZE_FMA</span></div>
<div class="line"><a name="l01665"></a><span class="lineno"> 1665</span>&#160;    <span class="comment">// Bandwidth of all of VFMA/MULPS/ADDPS is 0.5 on latest Intel processors.</span></div>
<div class="line"><a name="l01666"></a><span class="lineno"> 1666</span>&#160;    <span class="comment">// However for MULPS/ADDPS we have dependent sequence of 2 such</span></div>
<div class="line"><a name="l01667"></a><span class="lineno"> 1667</span>&#160;    <span class="comment">// instructions,</span></div>
<div class="line"><a name="l01668"></a><span class="lineno"> 1668</span>&#160;    <span class="comment">// so overall bandwidth is 1.0.</span></div>
<div class="line"><a name="l01669"></a><span class="lineno"> 1669</span>&#160;    <span class="keywordflow">if</span> (computeBandwidth == 0.5) computeBandwidth = 1.0;</div>
<div class="line"><a name="l01670"></a><span class="lineno"> 1670</span>&#160;<span class="preprocessor">#endif</span></div>
<div class="line"><a name="l01671"></a><span class="lineno"> 1671</span>&#160;    <span class="keywordflow">return</span> computeBandwidth;</div>
<div class="line"><a name="l01672"></a><span class="lineno"> 1672</span>&#160;  }</div>
<div class="line"><a name="l01673"></a><span class="lineno"> 1673</span>&#160; </div>
<div class="line"><a name="l01674"></a><span class="lineno"> 1674</span>&#160;};</div>
<div class="line"><a name="l01675"></a><span class="lineno"> 1675</span>&#160; </div>
<div class="line"><a name="l01676"></a><span class="lineno"> 1676</span>&#160;} <span class="comment">// end namespace Eigen</span></div>
<div class="line"><a name="l01677"></a><span class="lineno"> 1677</span>&#160; </div>
<div class="line"><a name="l01678"></a><span class="lineno"> 1678</span>&#160;<span class="preprocessor">#endif  </span><span class="comment">// EIGEN_USE_THREADS</span></div>
<div class="line"><a name="l01679"></a><span class="lineno"> 1679</span>&#160;<span class="preprocessor">#endif </span><span class="comment">// EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_THREAD_POOL_H</span></div>
<div class="ttc" id="agroup__Core__Module_html_ga0e45b6b2adead7c6a29815b99f9f14dd"><div class="ttname"><a href="../group__Core__Module.html#ga0e45b6b2adead7c6a29815b99f9f14dd">end</a></div><div class="ttdeci">static const lastp1_t end</div></div>
<div class="ttc" id="agroup__enums_html_gga45fe06e29902b7a2773de05ba27b47a1a4e19dd09d5ff42295ba1d72d12a46686"><div class="ttname"><a href="../group__enums.html#gga45fe06e29902b7a2773de05ba27b47a1a4e19dd09d5ff42295ba1d72d12a46686">Eigen::Unaligned</a></div><div class="ttdeci">Unaligned</div></div>
<div class="ttc" id="agroup__enums_html_ggaacded1a18ae58b0f554751f6cdf9eb13a0103672ae41005ab03b4176c765afd62"><div class="ttname"><a href="../group__enums.html#ggaacded1a18ae58b0f554751f6cdf9eb13a0103672ae41005ab03b4176c765afd62">Eigen::ColMajor</a></div><div class="ttdeci">ColMajor</div></div>
<div class="ttc" id="anamespaceEigen_html"><div class="ttname"><a href="namespaceEigen.html">Eigen</a></div><div class="ttdoc">Namespace containing all symbols from the Eigen library.</div></div>
<div class="ttc" id="anamespaceEigen_html_a62e77e0933482dafde8fe197d9a2cfde"><div class="ttname"><a href="../namespaceEigen.html#a62e77e0933482dafde8fe197d9a2cfde">Eigen::Index</a></div><div class="ttdeci">EIGEN_DEFAULT_DENSE_INDEX_TYPE Index</div></div>
</div><!-- fragment --></div><!-- contents -->
</div><!-- doc-content -->
<!-- start footer part -->
<div id="nav-path" class="navpath"><!-- id is needed for treeview function! -->
  <ul>
    <li class="navelem"><a class="el" href="dir_8d80078a9d0677ffcc4e95dc0faac368.html">unsupported</a></li><li class="navelem"><a class="el" href="dir_146d0efcf5acec477a223a5f4dae5578.html">Eigen</a></li><li class="navelem"><a class="el" href="dir_9f5a313d484c38c262f934ac7ef839a3.html">CXX11</a></li><li class="navelem"><a class="el" href="dir_26c37786afc5fa3d76cd91f33a83507e.html">src</a></li><li class="navelem"><a class="el" href="dir_a133238b105ca1b44753abc50c661b5b.html">Tensor</a></li><li class="navelem"><b>TensorContractionThreadPool.h</b></li>
    <li class="footer">Generated on Thu Apr 21 2022 13:07:59 for Eigen-unsupported by
    <a href="http://www.doxygen.org/index.html">
    <img class="footer" src="doxygen.png" alt="doxygen"/></a> 1.9.1 </li>
  </ul>
</div>
</body>
</html>
